finetune.py

# TODO: Add whisper
# TODO: Add BYOL-a /BYOL-s

import os
import json
import argparse
import datetime
from tqdm import tqdm
from time import sleep
from copy import deepcopy
from tqdm.contrib.logging import logging_redirect_tqdm
from utils import DotDict, NoFmtLog, get_logger, log_levels
from models import GE2E, Dense, ProbingModel, ProbingDense, TransformerBridge
from datasets import FeatureExtractorDataset, GPUDiskModeClassifierDataset, \
    CPUMemoryModeClassifierDataset, GPUMemoryModeClassifierDataset

import torch
import torchaudio
from torch.utils.data import DataLoader
from torcheval.metrics import MulticlassAccuracy


VERSION = '1.5'

FX_MODELS = ['WAV2VEC2_BASE','WAV2VEC2_LARGE',
            'WAV2VEC2_LARGE_XLSR','WAV2VEC2_LARGE_XLSR300M',
            'HUBERT_BASE', 'HUBERT_LARGE', 
            'HUBERT_ASR_LARGE', 'WAV2VEC2_ASR_LARGE_960H']
        
FX_MODEL_MAP = {'WAV2VEC2_BASE':'WAV2VEC2_BASE','WAV2VEC2_LARGE':'WAV2VEC2_LARGE',
                'WAV2VEC2_LARGE_XLSR':'WAV2VEC2_XLSR53','WAV2VEC2_LARGE_XLSR300M':'WAV2VEC2_XLSR_300M',
                'HUBERT_BASE':'HUBERT_BASE', 'HUBERT_LARGE':'HUBERT_LARGE',
                'HUBERT_ASR_LARGE':'HUBERT_ASR_LARGE', 'WAV2VEC2_ASR_LARGE_960H':'WAV2VEC2_ASR_LARGE_960H'}

CLF_MODELS = ['DENSE','PROBING','PROBING_DENSE']

DATASETS = ['AESDD','CaFE','EmoDB','EMOVO','IEMOCAP','RAVDESS','ShEMO']


class Trainer:

    def __init__(self, config):

        self.config = config

        # Setup Directories
        self.data_dir = os.path.join(self.config.data_dir, 'Audios', self.config.dataset)
        self.feature_dir = os.path.join(self.config.data_dir, 'Features', self.config.dataset, self.config.fx_model)
        self.history_dir = os.path.join(self.config.history_dir, f'v{VERSION}',self.config.dataset,
            f'FX_{self.config.fx_model}_CLF_{self.config.clf_model}', self.config.run_name)
        self.weights_dir = os.path.join(self.config.weights_dir, f'v{VERSION}',self.config.dataset,
            f'FX_{self.config.fx_model}_CLF_{self.config.clf_model}', self.config.run_name)

        if self.config.extract_mode == 'disk': os.makedirs(self.feature_dir, exist_ok=True)
        os.makedirs(self.history_dir, exist_ok=True)
        os.makedirs(self.weights_dir, exist_ok=True)

        # Setup logger
        log_file = os.path.join(self.history_dir, f'std.log')
        if os.path.isfile(log_file) : os.system(f'rm {log_file}')
        self.logger, self.no_fmt_logger = get_logger(self.config.run_name, self.config.log_level, log_file)
        self.no_fmt_log = NoFmtLog(self.no_fmt_logger)

        # Log configs
        self._print_banner()

        # Config check
        self._config_check()
    
        # Setup device
        if self.config.device == 'gpu': 
            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        else: self.device = 'cpu'
        self.logger.info(f'Running on {str(self.device).upper()}')
        self.no_fmt_log()

        # Load data info
        self.dataset_info = DotDict(json.load(open(os.path.join(self.data_dir,'info.json'))))
        
        # Setup models
        self.fx_model = self._get_feature_extractor()
        self.clf_model = self._get_classifier()
        if self.config.finetune_layer != None: self.ft_model = TransformerBridge(self.fx_model,self.clf_model,self.config.finetune_layer)
        
        # Setup metric
        self.acc_metric = MulticlassAccuracy(device=self.device)
        
        # Setup history
        self.history = DotDict(train_loss=[],train_acc=[],val_loss=[],val_acc=[],test_loss=[],test_acc=[])
        

    def _print_banner(self):
        try: console_width, _ = os.get_terminal_size(0)
        except : console_width = 50

        banner = '\n'+('-'*console_width)+f'\nMultilingual SER System v{VERSION}\n'+ ('-'*console_width) + '\n'
        max_len = max([len(key) for key in self.config.__dict__.keys()])+2
        for arg,val in self.config.__dict__.items():
            banner += ' '.join(arg.split('_')).title() + ' '*(max_len-len(arg))+': '+ str(val) + '\n'
        try: banner += 'Job Id' + ' '*(max_len-len('Job Id'))+': '+ str(os.environ['SLURM_JOB_ID']) + '\n'
        except: pass
        if torch.cuda.is_available() and self.config.device == 'gpu':
            banner += 'GPU' + ' '*(max_len-len('GPU'))+': '+ torch.cuda.get_device_name() + '\n'
            banner += 'GPU Count' + ' '*(max_len-len('GPU Count'))+': '+ str(torch.cuda.device_count()) + '\n'
            banner += 'GPU Memory' + ' '*(max_len-len('GPU Memory'))+': '+ str(round(torch.cuda.get_device_properties(0).total_memory/1024/1024/1024,2)) + ' GB\n'
        banner += f'\nTimestamp : {datetime.datetime.now()}\n'
        self.no_fmt_log(msg=banner)


    def _config_check(self):
        # TODO check for finetuning possiblity with given fx model
        # TODO check finetune layer if finetuning is possible layer less than 12 or 24
        # TODO check for possible extract modes implimented with a fx and clf
        pass


    def _get_feature_extractor(self):
        # Load Model
        try:
            if self.config.fx_model == 'GE2E':
                model = GE2E()
                checkpoint = torch.load(os.path.join(self.config.weights_dir,'pretrained',self.config.fx_model+'_weights.pt'), 'cpu')
                model.load_state_dict(checkpoint["model_state"])
            else:
                bundle = getattr(torchaudio.pipelines, FX_MODEL_MAP[self.config.fx_model])
                model = bundle.get_model()
        except:
            self.logger.exception(f"Could not load {self.config.fx_model} feature extractor")
            exit(0)
        return model


    def _get_classifier(self):
        
        if self.config.clf_model == 'PROBING': # orself.config.fx_model == 'GE2E'
            model = ProbingModel(self.config.fx_model, len(self.dataset_info.label_map))
        elif self.config.clf_model == 'PROBING_DENSE':
            model = ProbingDense(self.config.fx_model, len(self.dataset_info.label_map), self.device)
        elif self.config.clf_model == 'DENSE':
            model = Dense(self.config.fx_model, len(self.dataset_info.label_map), self.device)
           
        return model


    def _get_fx_dataloaders(self):
        test_dataset = FeatureExtractorDataset(self.config.fx_model, self.data_dir, 'test')
        train_dataset = FeatureExtractorDataset(self.config.fx_model, self.data_dir, 'train')
        val_dataset = FeatureExtractorDataset(self.config.fx_model, self.data_dir, 'validation')
        
        train_batch_size = 32
        test_batch_size = 32
        num_workers = self.config.num_workers

        test_loader = DataLoader(test_dataset, batch_size=test_batch_size, shuffle=False,
            num_workers=num_workers, collate_fn = test_dataset.data_collator)
        train_loader = DataLoader(train_dataset,batch_size=train_batch_size, shuffle=False,
            num_workers=num_workers, collate_fn = train_dataset.data_collator)
        val_loader = DataLoader(val_dataset, batch_size=test_batch_size, shuffle=False,
            num_workers=num_workers, collate_fn = val_dataset.data_collator)

        return DotDict(train=train_loader, test=test_loader, validation=val_loader)

    
    def _get_clf_dataloaders(self):
        
        if self.config.extract_mode=='gpu_disk':
            test_dataset = GPUDiskModeClassifierDataset(self.config.fx_model, self.feature_dir, self.dataset_info, 'test')
            train_dataset = GPUDiskModeClassifierDataset(self.config.fx_model, self.feature_dir, self.dataset_info, 'train')
            val_dataset = GPUDiskModeClassifierDataset(self.config.fx_model, self.feature_dir, self.dataset_info, 'validation')
        elif self.config.extract_mode=='cpu_memory':
            test_dataset = CPUMemoryModeClassifierDataset(self.config.fx_model, self.fx_model, self.data_dir, self.dataset_info, 'test')
            train_dataset = CPUMemoryModeClassifierDataset(self.config.fx_model, self.fx_model, self.data_dir, self.dataset_info, 'train')
            val_dataset = CPUMemoryModeClassifierDataset(self.config.fx_model, self.fx_model, self.data_dir, self.dataset_info, 'validation')
        elif self.config.extract_mode=='gpu_memory':
            test_dataset = GPUMemoryModeClassifierDataset(self.config.fx_model, self.data_dir, self.dataset_info, 'test')
            train_dataset = GPUMemoryModeClassifierDataset(self.config.fx_model, self.data_dir, self.dataset_info, 'train')
            val_dataset = GPUMemoryModeClassifierDataset(self.config.fx_model, self.data_dir, self.dataset_info, 'validation')

        train_batch_size = 32
        test_batch_size = 32
        num_workers = self.config.num_workers

        test_loader = DataLoader(test_dataset, batch_size=test_batch_size, shuffle=False,
            num_workers=num_workers, collate_fn = test_dataset.data_collator)
        train_loader = DataLoader(train_dataset,batch_size=train_batch_size, shuffle=True,
            num_workers=num_workers, collate_fn = train_dataset.data_collator)
        val_loader = DataLoader(val_dataset, batch_size=test_batch_size, shuffle=False,
            num_workers=num_workers, collate_fn = val_dataset.data_collator)

        return DotDict(train=train_loader, test=test_loader, validation=val_loader)


    def _train(self, dataloader, optimizer, criterion, progress_bar):
        total_loss = 0.0
        self.clf_model.train()
        for batch in dataloader:
            batch = tuple(input.to(self.device) for input in batch)

            optimizer.zero_grad()
            output = self.clf_model(batch[0],batch[1])

            loss = criterion(output, batch[2])
            
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
            self.acc_metric.update(output, batch[2])

            progress_bar.update(1)
        
        total_loss = total_loss / len(dataloader)
        accuracy = self.acc_metric.compute().item()*100

        self.acc_metric.reset()
        return total_loss, accuracy

    
    def _test(self, dataloader, criterion, progress_bar):
        total_loss = 0.0
        self.clf_model.eval()
        for batch in dataloader:
            batch = tuple(input.to(self.device) for input in batch)
           
            with torch.inference_mode():
                output = self.clf_model(batch[0],batch[1])
                loss = criterion(output, batch[2])

            total_loss += loss.item()
            self.acc_metric.update(output, batch[2])
            
            progress_bar.update(1)
        
        total_loss = total_loss / len(dataloader)
        accuracy = self.acc_metric.compute().item()*100
    
        self.acc_metric.reset()

        return total_loss, accuracy


    def _gpu_train(self, dataloader, optimizer, criterion, progress_bar):
        total_loss = 0.0
        self.clf_model.train()
        self.fx_model.eval()
        for batch in dataloader:
            batch = tuple(input.to(self.device) for input in batch)
            
            with torch.inference_mode():
                if self.config.fx_model == 'GE2E':
                    features = self.fx_model(batch[0])
                else:
                    features, _ = self.fx_model.extract_features(batch[0])
                    features = torch.stack([*features],dim=1)
            
            features = torch.clone(features)

            optimizer.zero_grad()
            output = self.clf_model(features,batch[1])

            loss = criterion(output, batch[2])
            
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
            self.acc_metric.update(output, batch[2])

            progress_bar.update(1)
        
        progress_bar.refresh()

        total_loss = total_loss / len(dataloader)
        accuracy = self.acc_metric.compute().item()*100

        self.acc_metric.reset()
        return total_loss, accuracy


    def _gpu_test(self, dataloader, criterion, progress_bar):
        total_loss = 0.0
        self.clf_model.eval()
        self.fx_model.eval()
        for batch in dataloader:
            batch = tuple(input.to(self.device) for input in batch)
           
            with torch.inference_mode():
                if self.config.fx_model == 'GE2E':
                    features = self.fx_model(batch[0])
                else:
                    features, _ = self.fx_model.extract_features(batch[0])
                    features = torch.stack([*features],dim=1)

                output = self.clf_model(features,batch[1])
                loss = criterion(output, batch[2])

            total_loss += loss.item()
            self.acc_metric.update(output, batch[2])
            
            progress_bar.update(1)
        
        progress_bar.refresh()

        total_loss = total_loss / len(dataloader)
        accuracy = self.acc_metric.compute().item()*100
    
        self.acc_metric.reset()

        return total_loss, accuracy


    def _gpu_probing_train(self, models, dataloader, optimizers, criterion, progress_bar):
        num_layers = len(models)
        total_losses = [0.0]*num_layers
        acc_metrics = []

        for layer in range(num_layers): 
            models[layer].train()
            metric = deepcopy(self.acc_metric)
            acc_metrics.append(metric)

        self.pos_conv = None
        def custom_hook(module, input_, output):  self.pos_conv = output
        try: handle = self.fx_model.encoder.transformer.pos_conv_embed.register_forward_hook(custom_hook)
        except: handle = self.fx_model.model.encoder.transformer.pos_conv_embed.register_forward_hook(custom_hook)

        self.fx_model.eval()
        for batch in dataloader:
            batch = tuple(input.to(self.device) for input in batch)
            
            with torch.inference_mode():
                features, _ = self.fx_model.extract_features(batch[0])
                features = torch.stack([self.pos_conv]+[*features],dim=1)

            features = torch.clone(features)

            for layer in range(num_layers):
                optimizers[layer].zero_grad()
                output = models[layer](features,batch[1],layer)

                loss = criterion(output, batch[2])
                
                loss.backward()
                optimizers[layer].step()
                
                total_losses[layer] += loss.item()
                acc_metrics[layer].update(output, batch[2])

                progress_bar.update(1)
        
        progress_bar.refresh()

        accuracies = []
        for layer in range(num_layers):
            total_losses[layer] = total_losses[layer] / len(dataloader)
            accuracies.append(acc_metrics[layer].compute().item()*100) 

        handle.remove()

        self.acc_metric.reset()
        return total_losses, accuracies


    def _gpu_probing_test(self, models, dataloader, criterion, progress_bar):
        num_layers = len(models)
        total_losses = [0.0]*num_layers
        acc_metrics = []

        for layer in range(num_layers): 
            models[layer].eval()
            metric = deepcopy(self.acc_metric)
            acc_metrics.append(metric)

        self.pos_conv = None
        def custom_hook(module, input_, output):  self.pos_conv = output
        try: handle = self.fx_model.encoder.transformer.pos_conv_embed.register_forward_hook(custom_hook)
        except: handle = self.fx_model.model.encoder.transformer.pos_conv_embed.register_forward_hook(custom_hook)

        self.fx_model.eval()
        for batch in dataloader:
            batch = tuple(input.to(self.device) for input in batch)
            
            with torch.inference_mode():
                features, _ = self.fx_model.extract_features(batch[0])
                features = torch.stack([self.pos_conv]+[*features],dim=1)

                for layer in range(num_layers):
                    output = models[layer](features,batch[1],layer)
                    loss = criterion(output, batch[2])
                    
                    total_losses[layer] += loss.item()
                    acc_metrics[layer].update(output, batch[2])

                    progress_bar.update(1)
        
        progress_bar.refresh()

        accuracies = []
        for layer in range(num_layers):
            total_losses[layer] = total_losses[layer] / len(dataloader)
            accuracies.append(acc_metrics[layer].compute().item()*100) 

        handle.remove()

        self.acc_metric.reset()
        return total_losses, accuracies


    def _gpu_finetune(self, model, dataloader, optimizer, criterion, progress_bar):
        total_loss = 0.0
        model.train()
        for batch in dataloader:
            batch = tuple(input.to(self.device) for input in batch)

            optimizer.zero_grad()
            output = self.clf_model(batch[0],batch[1])

            loss = criterion(output, batch[2])
            
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
            self.acc_metric.update(output, batch[2])

            progress_bar.update(1)
        
        progress_bar.refresh()

        total_loss = total_loss / len(dataloader)
        accuracy = self.acc_metric.compute().item()*100

        self.acc_metric.reset()
        return total_loss, accuracy


    def run_pipeline(self):
        
        if self.config.extract_mode=='gpu_disk':
            # Extract Features
            self.fx_dataloaders = self._get_fx_dataloaders()
            self._extract_features()
        
        # Setup classifier dataloaders
        self.clf_dataloaders = self._get_clf_dataloaders()

        # Setup criterion
        criterion = torch.nn.CrossEntropyLoss()

        # Set save paths
        best_model_path = os.path.join(self.weights_dir,'best_state.pt')

        train_pbar = tqdm(desc='Training   ', unit=' batch', colour='#EF5350', total=len(self.clf_dataloaders.train))
        val_pbar   = tqdm(desc='Validation ', unit=' batch', colour='#E0E0E0', total=len(self.clf_dataloaders.validation))
        test_pbar  = tqdm(desc='Testing    ', unit=' batch', colour='#42A5F5', total=len(self.clf_dataloaders.test))
        epoch_pbar = tqdm(desc='Epoch      ', unit=' epoch', colour='#43A047', total=self.config.epochs)

        # Train classifier
        self.clf_model.to(self.device)
        if self.config.extract_mode=='gpu_memory': self.fx_model.to(self.device)
        
        with logging_redirect_tqdm(loggers=[self.logger, self.no_fmt_logger]):
            self.no_fmt_log()
            self.logger.info('Training Classifier')
            self.no_fmt_log()

            if self.config.finetune_layer != None: 
                

            elif self.config.clf_model=='PROBING' or self.config.clf_model=='PROBING_DENSE': #and self.config.fx_model!='GE2E':
                models, optimizers = [], []
                if 'BASE' in self.config.fx_model: num_layers = 13
                elif 'LARGE' in self.config.fx_model:  num_layers = 25
                for _ in range(num_layers):
                    model = deepcopy(self.clf_model)
                    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
                    models.append(model)
                    optimizers.append(optimizer)

                train_pbar.unit = ' step'
                val_pbar.unit = ' step'
                test_pbar.unit = ' step'
                train_pbar.total *= num_layers
                val_pbar.total *= num_layers
                test_pbar.total *= num_layers
                train_pbar.refresh()
                val_pbar.refresh()
                test_pbar.refresh()
                best_model = DotDict(val_acc=[-1]*num_layers, test_acc=[0]*num_layers)
            else:
                optimizer = torch.optim.Adam(self.clf_model.parameters(), lr=0.001)
                best_model = DotDict(val_acc=-1, test_acc=0)
                

            for epoch in range(1,self.config.epochs+1):
                
                if self.config.clf_model=='PROBING' or self.config.clf_model=='PROBING_DENSE': #and self.config.fx_model!='GE2E':
                    train_loss, train_acc = self._gpu_probing_train(models, self.clf_dataloaders.train, optimizers, criterion, train_pbar)
                    val_loss, val_acc = self._gpu_probing_test(models, self.clf_dataloaders.validation, criterion, val_pbar)
                    test_loss, test_acc = self._gpu_probing_test(models, self.clf_dataloaders.test, criterion, test_pbar)
                    
                    epoch_pbar.update(1)

                    for layer in range(num_layers):
                        pretty_space = ' '*(2-len(str(layer+1)))
                        self.logger.info('Epoch: %s  |  Layer: %s%s  |  Train Loss: %.3f | Train Acc: %.2f  |  Val Loss: %.3f | Val Acc: %.2f  |  Test Loss: %.3f | Tes Acc: %.2f' \
                                    %(epoch, pretty_space, layer+1, train_loss[layer], train_acc[layer], val_loss[layer], val_acc[layer], test_loss[layer], test_acc[layer]))
                        if val_acc[layer] > best_model.val_acc[layer]:
                            best_model.val_acc[layer] = val_acc[layer]
                            best_model.test_acc[layer] = test_acc[layer]
                    self.no_fmt_log()                   

                else:
                    if self.config.extract_mode=='gpu_memory':
                        train_loss, train_acc = self._gpu_train(self.clf_dataloaders.train, optimizer, criterion, train_pbar)
                        val_loss, val_acc = self._gpu_test(self.clf_dataloaders.validation, criterion, val_pbar)
                        test_loss, test_acc = self._gpu_test(self.clf_dataloaders.test, criterion, test_pbar)
                    elif self.config.extract_mode=='cpu_memory' or self.config.extract_mode=='gpu_disk':
                        train_loss, train_acc = self._train(self.clf_dataloaders.train, optimizer, criterion, train_pbar)
                        val_loss, val_acc = self._test(self.clf_dataloaders.validation, criterion, val_pbar)
                        test_loss, test_acc = self._test(self.clf_dataloaders.test, criterion, test_pbar)

                    epoch_pbar.update(1)
                    self.logger.info('Epoch: %s  |  Train Loss: %.3f | Train Acc: %.2f  |  Val Loss: %.3f | Val Acc: %.2f  |  Test Loss: %.3f | Tes Acc: %.2f' \
                                    %(epoch, train_loss, train_acc, val_loss, val_acc, test_loss, test_acc))

                    if val_acc > best_model.val_acc:
                        best_model.val_acc = val_acc
                        best_model.test_acc = test_acc
                        state = dict(
                            epoch=epoch,
                            val_acc=val_acc,
                            test_acc=test_acc,
                            model_state=self.clf_model.state_dict(),
                            optimizer_state=optimizer.state_dict())
                        torch.save(state,best_model_path)
                        self.logger.debug(f'Best state saved | Val Acc {val_acc} | Test Acc {test_acc}')
                        self.no_fmt_log(level='debug')

                self.history.train_loss.append(train_loss)
                self.history.train_acc.append(train_acc)
                self.history.val_loss.append(val_loss)
                self.history.val_acc.append(val_acc)
                self.history.test_loss.append(test_loss)
                self.history.test_acc.append(test_acc)

                train_pbar.reset()
                val_pbar.reset()
                test_pbar.reset()

            train_pbar.update(train_pbar.total)
            val_pbar.update(val_pbar.total)
            test_pbar.update(test_pbar.total)
            self.no_fmt_log()
            time_taken = epoch_pbar.format_dict['elapsed']
            train_pbar.close()
            val_pbar.close()
            test_pbar.close()
            epoch_pbar.close()

        self.no_fmt_log()
        self.logger.info(f'Time Taken: {epoch_pbar.format_interval(time_taken)}')
        self.no_fmt_log()

        if self.config.clf_model=='PROBING' or self.config.clf_model=='PROBING_DENSE': #and self.config.fx_model!='GE2E':
            
            for layer in range(num_layers):
                pretty_space = ' '*(2-len(str(layer+1)))
                self.logger.info('Layer: %s%s  |  Best Val Acc: %.2f  |  Best Test Acc: %.2f' \
                                    %(pretty_space, layer+1, best_model.val_acc[layer], best_model.test_acc[layer]))
            self.no_fmt_log()

            # Save last state
            last_state_path = os.path.join(self.weights_dir, 'last_states.pt')
            states=[]
            for layer in range(num_layers):
                
                state = dict(
                    epoch=epoch,
                    val_acc=val_acc[layer],
                    test_acc=test_acc[layer],
                    model_state=models[layer].state_dict(),
                    optimizer_state=optimizers[layer].state_dict())
                states.append(state)
            torch.save(states,last_state_path)
            self.logger.debug(f'Last state saved')
        else:
            self.logger.info(f'Best Val Acc : {best_model.val_acc} | Best Test Acc : {best_model.test_acc}')
            self.no_fmt_log()

            # Agg weights
            if self.config.fx_model != 'GE2E' and self.config.clf_model == 'DENSE':
                agg_weights = ' '.join([str(weight[0]) for weight in self.clf_model.aggr.state_dict()['weight'][0].detach().cpu().tolist()])
                self.logger.info(f'Agg. Weights : \n{agg_weights}\n')
                agg_weight_path = os.path.join(self.history_dir, 'agg_weights.pt')
                torch.save(agg_weights,agg_weight_path)

            # Save last state
            last_state_path = os.path.join(self.weights_dir, 'last_state.pt')
            state = dict(
                epoch=epoch,
                val_acc=val_acc,
                test_acc=test_acc,
                model_state=self.clf_model.state_dict(),
                optimizer_state=optimizer.state_dict())
            torch.save(state,last_state_path)
            self.logger.debug(f'Last state saved | Val Acc {val_acc} | Test Acc {test_acc}')

        # Save history
        history_path = os.path.join(self.history_dir, 'history.pt')
        torch.save(dict(self.history),history_path)
        self.no_fmt_log()
        self.logger.info(f'History saved ({history_path})')

        return self.history


    def finetune_pipeline(self):
        
        # Setup classifier dataloaders
        self.clf_dataloaders = self._get_clf_dataloaders()

        # Setup criterion
        criterion = torch.nn.CrossEntropyLoss()

        # Set save paths
        best_model_path = os.path.join(self.weights_dir,'best_state.pt')

        train_pbar = tqdm(desc='Training   ', unit=' batch', colour='#EF5350', total=len(self.clf_dataloaders.train))
        val_pbar   = tqdm(desc='Validation ', unit=' batch', colour='#E0E0E0', total=len(self.clf_dataloaders.validation))
        test_pbar  = tqdm(desc='Testing    ', unit=' batch', colour='#42A5F5', total=len(self.clf_dataloaders.test))
        epoch_pbar = tqdm(desc='Epoch      ', unit=' epoch', colour='#43A047', total=self.config.epochs)

        # Train classifier
        self.clf_model.to(self.device)
        if self.config.extract_mode=='gpu_memory': self.fx_model.to(self.device)
        
        with logging_redirect_tqdm(loggers=[self.logger, self.no_fmt_logger]):
            self.no_fmt_log()
            self.logger.info('Training Classifier')
            self.no_fmt_log()

            if self.config.clf_model=='PROBING' or self.config.clf_model=='PROBING_DENSE': #and self.config.fx_model!='GE2E':
                models, optimizers = [], []
                if 'BASE' in self.config.fx_model: num_layers = 13
                elif 'LARGE' in self.config.fx_model:  num_layers = 25
                for _ in range(num_layers):
                    model = deepcopy(self.clf_model)
                    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
                    models.append(model)
                    optimizers.append(optimizer)

                train_pbar.unit = ' step'
                val_pbar.unit = ' step'
                test_pbar.unit = ' step'
                train_pbar.total *= num_layers
                val_pbar.total *= num_layers
                test_pbar.total *= num_layers
                train_pbar.refresh()
                val_pbar.refresh()
                test_pbar.refresh()
                best_model = DotDict(val_acc=[-1]*num_layers, test_acc=[0]*num_layers)
            else:
                optimizer = torch.optim.Adam(self.clf_model.parameters(), lr=0.001)
                best_model = DotDict(val_acc=-1, test_acc=0)
                

            for epoch in range(1,self.config.epochs+1):
                
                if self.config.clf_model=='PROBING' or self.config.clf_model=='PROBING_DENSE': #and self.config.fx_model!='GE2E':
                    train_loss, train_acc = self._gpu_probing_train(models, self.clf_dataloaders.train, optimizers, criterion, train_pbar)
                    val_loss, val_acc = self._gpu_probing_test(models, self.clf_dataloaders.validation, criterion, val_pbar)
                    test_loss, test_acc = self._gpu_probing_test(models, self.clf_dataloaders.test, criterion, test_pbar)
                    
                    epoch_pbar.update(1)

                    for layer in range(num_layers):
                        pretty_space = ' '*(2-len(str(layer+1)))
                        self.logger.info('Epoch: %s  |  Layer: %s%s  |  Train Loss: %.3f | Train Acc: %.2f  |  Val Loss: %.3f | Val Acc: %.2f  |  Test Loss: %.3f | Tes Acc: %.2f' \
                                    %(epoch, pretty_space, layer+1, train_loss[layer], train_acc[layer], val_loss[layer], val_acc[layer], test_loss[layer], test_acc[layer]))
                        if val_acc[layer] > best_model.val_acc[layer]:
                            best_model.val_acc[layer] = val_acc[layer]
                            best_model.test_acc[layer] = test_acc[layer]
                    self.no_fmt_log()                   

                else:
                    if self.config.extract_mode=='gpu_memory':
                        train_loss, train_acc = self._gpu_train(self.clf_dataloaders.train, optimizer, criterion, train_pbar)
                        val_loss, val_acc = self._gpu_test(self.clf_dataloaders.validation, criterion, val_pbar)
                        test_loss, test_acc = self._gpu_test(self.clf_dataloaders.test, criterion, test_pbar)
                    elif self.config.extract_mode=='cpu_memory' or self.config.extract_mode=='gpu_disk':
                        train_loss, train_acc = self._train(self.clf_dataloaders.train, optimizer, criterion, train_pbar)
                        val_loss, val_acc = self._test(self.clf_dataloaders.validation, criterion, val_pbar)
                        test_loss, test_acc = self._test(self.clf_dataloaders.test, criterion, test_pbar)

                    epoch_pbar.update(1)
                    self.logger.info('Epoch: %s  |  Train Loss: %.3f | Train Acc: %.2f  |  Val Loss: %.3f | Val Acc: %.2f  |  Test Loss: %.3f | Tes Acc: %.2f' \
                                    %(epoch, train_loss, train_acc, val_loss, val_acc, test_loss, test_acc))

                    if val_acc > best_model.val_acc:
                        best_model.val_acc = val_acc
                        best_model.test_acc = test_acc
                        state = dict(
                            epoch=epoch,
                            val_acc=val_acc,
                            test_acc=test_acc,
                            model_state=self.clf_model.state_dict(),
                            optimizer_state=optimizer.state_dict())
                        torch.save(state,best_model_path)
                        self.logger.debug(f'Best state saved | Val Acc {val_acc} | Test Acc {test_acc}')
                        self.no_fmt_log(level='debug')

                self.history.train_loss.append(train_loss)
                self.history.train_acc.append(train_acc)
                self.history.val_loss.append(val_loss)
                self.history.val_acc.append(val_acc)
                self.history.test_loss.append(test_loss)
                self.history.test_acc.append(test_acc)

                train_pbar.reset()
                val_pbar.reset()
                test_pbar.reset()

            train_pbar.update(train_pbar.total)
            val_pbar.update(val_pbar.total)
            test_pbar.update(test_pbar.total)
            self.no_fmt_log()
            time_taken = epoch_pbar.format_dict['elapsed']
            train_pbar.close()
            val_pbar.close()
            test_pbar.close()
            epoch_pbar.close()

        self.no_fmt_log()
        self.logger.info(f'Time Taken: {epoch_pbar.format_interval(time_taken)}')
        self.no_fmt_log()

        if self.config.clf_model=='PROBING' or self.config.clf_model=='PROBING_DENSE': #and self.config.fx_model!='GE2E':
            
            for layer in range(num_layers):
                pretty_space = ' '*(2-len(str(layer+1)))
                self.logger.info('Layer: %s%s  |  Best Val Acc: %.2f  |  Best Test Acc: %.2f' \
                                    %(pretty_space, layer+1, best_model.val_acc[layer], best_model.test_acc[layer]))
            self.no_fmt_log()

            # Save last state
            last_state_path = os.path.join(self.weights_dir, 'last_states.pt')
            states=[]
            for layer in range(num_layers):
                
                state = dict(
                    epoch=epoch,
                    val_acc=val_acc[layer],
                    test_acc=test_acc[layer],
                    model_state=models[layer].state_dict(),
                    optimizer_state=optimizers[layer].state_dict())
                states.append(state)
            torch.save(states,last_state_path)
            self.logger.debug(f'Last state saved')
        else:
            self.logger.info(f'Best Val Acc : {best_model.val_acc} | Best Test Acc : {best_model.test_acc}')
            self.no_fmt_log()

            # Agg weights
            if self.config.fx_model != 'GE2E' and self.config.clf_model == 'DENSE':
                agg_weights = ' '.join([str(weight[0]) for weight in self.clf_model.aggr.state_dict()['weight'][0].detach().cpu().tolist()])
                self.logger.info(f'Agg. Weights : \n{agg_weights}\n')
                agg_weight_path = os.path.join(self.history_dir, 'agg_weights.pt')
                torch.save(agg_weights,agg_weight_path)

            # Save last state
            last_state_path = os.path.join(self.weights_dir, 'last_state.pt')
            state = dict(
                epoch=epoch,
                val_acc=val_acc,
                test_acc=test_acc,
                model_state=self.clf_model.state_dict(),
                optimizer_state=optimizer.state_dict())
            torch.save(state,last_state_path)
            self.logger.debug(f'Last state saved | Val Acc {val_acc} | Test Acc {test_acc}')

        # Save history
        history_path = os.path.join(self.history_dir, 'history.pt')
        torch.save(dict(self.history),history_path)
        self.no_fmt_log()
        self.logger.info(f'History saved ({history_path})')

        return self.history

def get_args():
    """Parse input arguments"""
    parser = argparse.ArgumentParser(f'Multilingual SER System v{VERSION}')

    parser.add_argument("-r","--run_name", metavar="<str>", default="test", type=str,
                        help='Run Name') 
    parser.add_argument("-fm","--fx_model", metavar="<str>", default="GE2E", type=str,
                        choices=FX_MODELS, help=str(FX_MODELS))
    parser.add_argument("-cm","--clf_model", metavar="<str>", default="PROBING", type=str,
                        choices=CLF_MODELS, help=str(CLF_MODELS))
    parser.add_argument("-em","--extract_mode", metavar="<str>", default="cpu_memory", type=str,
                        choices=['gpu_disk','gpu_memory','cpu_memory'], help='GPU Disk mode will extract features on GPU and will save the features \
                            to disk and then training will continue using disk cache, GPU Memory mode will extract features and train the model both on GPU,\
                            CPU Memory mode will extract features on CPU and run while training on GPU') 
    parser.add_argument("-d","--dataset", metavar="<str>", default="EmoDB", type=str,
                        choices=DATASETS, help=str(DATASETS))


    parser.add_argument("-dv","--device", metavar="<str>", default="gpu", type=str,
                        choices=['cpu','gpu'], help='Device to run on') 
    parser.add_argument("-e","--epochs", metavar="<int>", default=20, type=int,
                        help='Number of training epochs')
    parser.add_argument("-nw","--num_workers", metavar="<int>", default=2, type=int,
                        help='Number of dataloader workers')


    parser.add_argument("-dd","--data_dir", metavar="<dir>", default="./data", type=str,
                        help='Data directory')     
    parser.add_argument("-hd","--history_dir", metavar="<dir>", default="./history", type=str,
                        help='History directory')     
    parser.add_argument("-wd","--weights_dir", metavar="<dir>", default="./weights", type=str,
                        help='Weights directory')


    parser.add_argument("-ll","--log_level", metavar="<str>", default="info", type=str,
                        choices=list(log_levels.keys()), help=str(list(log_levels.keys())))
    parser.add_argument("-pc","--purge_cache", action="store_true", default=False,
                        help='Purge cached features and extract them again')  
    parser.add_argument("-jn","--job_name", metavar='<str>', default='Manual_Run', type=str,
                        help='SLURM Job Name for logging')                                            

    return parser.parse_args()


if __name__ == '__main__':
    args = get_args()
        
    # Test
    if args.run_name =='test':

        args.dataset = 'EMOVO'
        # args.fx_model = 'HUBERT_BASE'
        #' ['WAV2VEC2_BASE','WAV2VEC2_LARGE','WAV2VEC2_LARGE_XLSR','WAV2VEC2_LARGE_XLSR300M','HUBERT_BASE','HUBERT_LARGE','WAV2VEC2_ASR_LARGE_960H', 'HUBERT_ASR_LARGE']
        args.fx_model = 'WAV2VEC2_ASR_LARGE_960H'
        args.clf_model = 'DENSE'
        args.extract_mode ='gpu_memory'
        args.history_dir = './test/history'
        args.weights_dir = './test/weights'
        args.log_level='debug'
        args.num_workers=3

        args.epochs = 1

        os.system('rm -rf ./test')
        os.system('mkdir test')
        os.system('mkdir test/weights')
        os.system('ln -s /scratch/as14229/Projects/Multilingual-Speech-Emotion-Recognition-System/weights/pretrained test/weights/pretrained')

    # Train Classifier
    trainer = Trainer(args)
    trainer.train_pipeline()