trainer.py

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import pytorch_lightning as pl
from sklearn import metrics
from transformers import AutoModelForAudioClassification
import numpy as np
from collections import OrderedDict
from torchmetrics import MeanMetric, MaxMetric, Accuracy
import torchmetrics.functional as tmf

from model.linear import FeedforwardModel
from model.linear_small import FeedforwardModelSmall
from model.linear_attn_ck import FeedforwardModelAttnCK
from model.linear_mt import FeedforwardModelMT
from model.linear_mt_attn_ck import FeedforwardModelMTAttnCK

import logging
import yaml
from omegaconf import DictConfig

import torch
from torch.distributed import all_gather, get_world_size
# from lion_pytorch import Lion
from torch_optimizer import RAdam

def gather_all_results(tensor):
    """
    Gather tensors from all GPUs in distributed training.
    """
    gathered_tensors = [torch.zeros_like(tensor) for _ in range(get_world_size())]
    all_gather(gathered_tensors, tensor)
    return torch.cat(gathered_tensors, dim=0)

# torch.set_float32_matmul_precision('medium')

log = logging.getLogger(__name__)
class MusicClassifier(pl.LightningModule):
    def __init__(self, cfg: DictConfig, output_file = None):
        super(MusicClassifier, self).__init__()
        self.cfg = cfg
        self.encoder = cfg.model.encoder
        self.classifier = cfg.model.classifier
        self.lr = cfg.model.lr
        self.output_file = output_file
        self.kd = cfg.model.kd
        self.kd_weight = cfg.model.kd_weight
        self.kd_temperature = self.cfg.model.kd_temperature

        layer_size = len(self.cfg.model.layers)
        mert_dim = 768 * layer_size

        self.feature_dim_dict = {
            "MERT": mert_dim
        }

        encoders = self.encoder.split("-")
        self.input_size = sum(self.feature_dim_dict[encoder] for encoder in encoders)
        self.num_datasets = len(self.cfg.datasets)

        if "mt" in self.classifier:
            if self.num_datasets < 2:
                raise Exception("Error: Dataset size >= 2 needed for MT classifier")
            classifiers = {
                "linear-mt-attn-ck": FeedforwardModelMTAttnCK,
            }
            if self.classifier in classifiers:
                self.model = classifiers[self.classifier](
                    input_size=self.input_size,
                    output_size_classification=56,
                    output_size_regression=2
                )
            else:
                raise Exception(f"Unknown classifier: {self.classifier}")          
        else:
            if self.num_datasets  >= 2:
                raise Exception(f"Error: Dataset size == 1 needed for classifier") 
            dataset_name = self.cfg.datasets[0]
            self.output_size = self.cfg.dataset[dataset_name].output_size
            classifiers = {
                "linear": FeedforwardModel,
                "linear-attn-ck": FeedforwardModelAttnCK
            }

            if self.classifier in classifiers:
                self.model = classifiers[self.classifier](input_size=self.input_size, output_size=self.output_size)
            else:
                raise Exception(f"Unknown classifier: {self.classifier}")
        
                
        if self.kd:
            self.teacher_models = {}
            
            for dataset in self.cfg.datasets:
                self.output_size = self.cfg.dataset[dataset].output_size
                teacher_model_path = getattr(self.cfg, f"checkpoint_{dataset}", None)
                
                if teacher_model_path:
                    # Create a new teacher model instance
                    teacher_model = FeedforwardModelAttnCK(
                        input_size=self.input_size,
                        output_size=self.output_size,
                    )
                    
                    # Load the checkpoint
                    checkpoint = torch.load(teacher_model_path, map_location=self.device, weights_only=False)
                    state_dict = checkpoint["state_dict"]
                    
                    # Adjust the keys in the state_dict
                    state_dict = {key.replace("model.", ""): value for key, value in state_dict.items()}
                    
                    # Filter state_dict to match model's keys
                    model_keys = set(teacher_model.state_dict().keys())
                    filtered_state_dict = {key: value for key, value in state_dict.items() if key in model_keys}
                    
                    # Load the filtered state_dict and set the model to evaluation mode
                    teacher_model.load_state_dict(filtered_state_dict)
                    teacher_model.to(self.device)

                    teacher_model.eval()
                    
                    # Store the teacher model in the dictionary with the dataset name as the key
                    self.teacher_models[dataset] = teacher_model

        probas = torch.from_numpy(np.load("dataset/jamendo/meta/probas_train.npy"))
        pos_weight = torch.tensor(1.) / probas
        weight = torch.tensor(2.) / (torch.tensor(1.) + pos_weight)
        
        self.loss_fn_classification = nn.BCEWithLogitsLoss(
            pos_weight=pos_weight,reduction="mean",weight=weight
        )
        self.loss_fn_classification_eval = nn.BCEWithLogitsLoss(
            pos_weight=pos_weight,reduction="none",weight=weight
        )
        
        self.loss_fn_regression = nn.MSELoss()  

        self.loss_kd = nn.KLDivLoss(reduction="batchmean")

        self.prd_array = []
        self.gt_array = []
        self.song_array = []

        self.prd_array_va = []
        self.gt_array_va = []
        self.song_array_va = []

        self.validation_predictions = []
        self.validation_targets = []
        self.validation_results = {'preds': [], 'gt': []}

        self.trn_loss = MeanMetric()
        self.val_loss = MeanMetric()

    def forward(self, model_input_dic, output_idx = 0):
        if "mt" in self.classifier:
            classification_output, regression_output = self.model(model_input_dic)
            if output_idx == 0:
                return classification_output
            elif output_idx == 1:
                return regression_output
            elif output_idx == 2:
                return classification_output, regression_output
        else:
            output = self.model(model_input_dic)
            return output
        
    def compute_classification_loss(self, model_input_dic, y_mood):
        classification_logits = self(model_input_dic, 0)
        loss= self.loss_fn_classification(classification_logits, y_mood)
        return loss

    def compute_regression_loss(self, model_input_dic, y_va):
        regression_output = self(model_input_dic, 1)
        loss = self.loss_fn_regression(regression_output, y_va)
        return loss
    
    def compute_mt_loss(self, model_input_dic, y_mood, y_va):
        classification_logits, regression_output = self(model_input_dic, 2)
        loss_classification = self.loss_fn_classification(classification_logits, y_mood)
        loss_regression = self.loss_fn_regression(regression_output, y_va)
        return loss_classification, loss_regression
    

    def compute_kd_loss(self, model_input_dic, y_mood, y_va, dataset_name):
        """
        Compute knowledge distillation loss for a given dataset.
        """
        # Forward pass through student model
        s_logits_mood, s_logits_va = self(model_input_dic, 2)

        # Compute student losses
        s_loss_mood = self.loss_fn_classification(s_logits_mood, y_mood)
        s_loss_va = self.loss_fn_regression(s_logits_va, y_va)

        # Get the corresponding teacher model for the dataset
        teacher_model = self.teacher_models.get(dataset_name)
        teacher_model.to(self.device)

        # Ensure teacher model exists
        if teacher_model is None:
            raise ValueError(f"No teacher model found for dataset: {dataset_name}")

        with torch.no_grad():
            # Forward pass through teacher model
            t_logits = teacher_model(model_input_dic)

        # Compute knowledge distillation losses
        t_probs = torch.softmax(t_logits / self.kd_temperature, dim=-1)
        if dataset_name == "jamendo":
            s_probs_mood = torch.log_softmax(s_logits_mood / self.kd_temperature, dim=-1)
            kd_loss = self.loss_kd(s_probs_mood, t_probs)
        else:
            s_probs_va = torch.log_softmax(s_logits_va / self.kd_temperature, dim=-1)
            kd_loss = self.loss_kd(s_probs_va, t_probs)

        return kd_loss, s_loss_mood, s_loss_va
    
    def handle_dataset(self, dataset_name, batch, losses, total_loss, stage):
        dataset_batch = batch[dataset_name]

        model_input_dic = {}
        model_input_dic["x_mert"] = dataset_batch["x_mert"]
        model_input_dic["x_chord"] = dataset_batch["x_chord"]
        model_input_dic["x_chord_root"] = dataset_batch["x_chord_root"]
        model_input_dic["x_chord_attr"] = dataset_batch["x_chord_attr"]
        model_input_dic["x_key"] = dataset_batch["x_key"]

        if "mt" in self.classifier:            
            if dataset_name == "jamendo":
                y_mood = dataset_batch["y_mood"]
                y_va = dataset_batch["y_va"]
                if self.kd:
                    kd_loss, s_loss_mood, s_loss_va = self.compute_kd_loss(model_input_dic, y_mood, y_va, dataset_name)
                    if stage == "train":
                        losses['loss_mood'] = s_loss_mood

                        total_loss += self.kd_weight * kd_loss + (1 - self.kd_weight) * s_loss_mood
                    else:
                        losses['loss_mood'] = s_loss_mood
                        total_loss += s_loss_mood
                else:
                    s_loss_mood, s_loss_va = self.compute_mt_loss(model_input_dic, y_mood, y_va)
                    if stage == "train":
                        losses['loss_mood'] = s_loss_mood
                        total_loss += s_loss_mood 
                    else:
                        losses['loss_mood'] = s_loss_mood
                        total_loss += s_loss_mood
            else:
                y_mood = dataset_batch["y_mood"]
                y_va = dataset_batch["y_va"]

                if self.kd:
                    kd_loss, s_loss_mood, s_loss_va = self.compute_kd_loss(model_input_dic, y_mood, y_va, dataset_name)
                    if stage == "train":
                        losses['loss_va'] = s_loss_va
                        total_loss += self.kd_weight * kd_loss + (1 - self.kd_weight) * s_loss_va
                    else:
                        losses['loss_va'] = s_loss_va
                        total_loss += s_loss_va
                else:
                    s_loss_mood, s_loss_va = self.compute_mt_loss(model_input_dic, y_mood, y_va)
                    if stage == "train":
                        losses['loss_va'] = s_loss_va
                        total_loss += s_loss_va
                    else:
                        losses['loss_va'] = s_loss_va
                        total_loss += s_loss_va
        else:
            if dataset_name == "jamendo":
                y_mood = dataset_batch["y_mood"]
                loss_classification = self.compute_classification_loss(model_input_dic, y_mood)
                losses['loss_mood'] = loss_classification
                total_loss += loss_classification
            else:
                y_va = dataset_batch["y_va"]
                loss_regression = self.compute_regression_loss(model_input_dic, y_va)
                losses['loss_va'] = loss_regression
                total_loss += loss_regression

        return total_loss

    def training_step(self, batch, batch_idx):
        total_loss = 0
        losses = {}
        datasets = ["jamendo",  "deam", "emomusic", "pmemo"]
        
        for dataset in datasets:
            if dataset in batch and batch[dataset] is not None:
                total_loss = self.handle_dataset(dataset, batch, losses, total_loss, "train")
        
        batch_size = batch[next(iter(batch))]["x_mert"].size(0)

        self.log('train_loss_mood', losses.get('loss_mood', 0), on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True, batch_size=batch_size)
        self.log('train_loss_va', losses.get('loss_va', 0), on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True, batch_size=batch_size)
        self.log('train_loss', total_loss, on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True, batch_size=batch_size)

        return total_loss
    
    def validation_step(self, batch, batch_idx):

        total_loss = 0
        losses = {}
        datasets = ["jamendo", "deam", "emomusic", "pmemo"]

        for dataset in datasets:
            if dataset in batch and batch[dataset] is not None:
                total_loss = self.handle_dataset(dataset, batch, losses, total_loss, "val")
        
        batch_size = batch[next(iter(batch))]["x_mert"].size(0)

        self.log('val_loss_mood', losses.get('loss_mood', 0), on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True, batch_size=batch_size)
        self.log('val_loss_va', losses.get('loss_va', 0), on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True, batch_size=batch_size)
        self.log('val_loss', total_loss, on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True, batch_size=batch_size)
        return total_loss

    def test_step(self, batch, batch_idx):
        total_loss = 0
        losses = {}
        datasets = ["jamendo",  "deam", "emomusic", "pmemo"]

        for dataset in datasets:
            if dataset in batch and batch[dataset] is not None:
                dataset_batch = batch[dataset]

                model_input_dic = {}
                model_input_dic["x_mert"] = dataset_batch["x_mert"]
                model_input_dic["x_chord"] = dataset_batch["x_chord"]
                model_input_dic["x_chord_root"] = dataset_batch["x_chord_root"]
                model_input_dic["x_chord_attr"] = dataset_batch["x_chord_attr"]
                
                model_input_dic["x_key"] = dataset_batch["x_key"]

                if dataset == "jamendo":
                    y_mood = dataset_batch["y_mood"]
                    classification_logits = self(model_input_dic, 0)
                    
                    loss_classification = self.loss_fn_classification(classification_logits, y_mood)
                    total_loss += loss_classification

                    probs = torch.sigmoid(classification_logits)
                    if not hasattr(self, 'jamendo_results'):
                        self.jamendo_results = {'preds': [], 'gt': [], 'paths': []}
                    
                    self.jamendo_results['preds'].extend(probs.detach().cpu().numpy())
                    self.jamendo_results['gt'].extend(y_mood.detach().cpu().numpy())
                    self.jamendo_results['paths'].extend(dataset_batch["path"])

                    losses['test_loss_mood'] = loss_classification

                else:  # Handle regression for all other datasets
                    if batch[dataset] is not None:
                        y_va = dataset_batch["y_va"]
                        regression_output = self(model_input_dic, 1)

                        loss_regression = self.loss_fn_regression(regression_output, y_va)
                        total_loss += loss_regression

                        # Track results separately for each dataset
                        if not hasattr(self, f'{dataset}_results'):
                            setattr(self, f'{dataset}_results', {'preds': [], 'gt': [], 'paths': []})
                        
                        dataset_results = getattr(self, f'{dataset}_results')
                        dataset_results['preds'].extend(regression_output.detach().cpu().numpy())
                        dataset_results['gt'].extend(y_va.detach().cpu().numpy())
                        dataset_results['paths'].extend(dataset_batch["path"])

                        losses['test_loss_va'] = loss_regression

        batch_size = batch[next(iter(batch))]["x_mert"].size(0)

        # Log the classification and regression losses
        self.log('test_loss_mood', losses.get('test_loss_mood', 0), on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True, batch_size=batch_size)
        self.log('test_loss_va', losses.get('test_loss_va', 0), on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True, batch_size=batch_size)
        self.log('test_loss', total_loss, on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True, batch_size=batch_size)

        return total_loss

    def on_test_end(self):
        output_dic = {}

        # Jamendo classification metrics (AUC and PR AUC)
        if hasattr(self, 'jamendo_results') and self.jamendo_results['preds']:
            roc_auc, pr_auc = self.get_auc(self.jamendo_results['preds'], self.jamendo_results['gt'])

            roc_auc = roc_auc.item()
            pr_auc = pr_auc.item()

            log.info('*** Display ROC_AUC_MACRO scores (Jamendo) ***')
            log.info(f"ROC_AUC_MACRO: {round(roc_auc, 4)}")
            log.info(f"PR_AUC_MACRO: {round(pr_auc, 4)}")

            if self.output_file is not None:
                with open(self.output_file, 'a') as f:
                    f.write(f"ROC_AUC_MACRO (Jamendo): {round(roc_auc, 4)}\n")
                    f.write(f"PR_AUC_MACRO (Jamendo): {round(pr_auc, 4)}\n")

            output_dic["test_roc_auc_jamendo"] = round(roc_auc, 4)
            output_dic["test_pr_auc_jamendo"] = round(pr_auc, 4)

        # Metrics for each regression dataset (DMDD, DEAM, EmoMusic, PMEmo)
        for dataset in ["deam", "emomusic", "pmemo"]:
            dataset_results = getattr(self, f'{dataset}_results', None)

            if dataset_results and dataset_results['preds']:
                preds = torch.tensor(np.array(dataset_results['preds']))
                gts = torch.tensor(np.array(dataset_results['gt']))

                # Assuming valence is the first column and arousal is the second
                preds_valence = preds[:, 0]
                preds_arousal = preds[:, 1]
                gts_valence = gts[:, 0]
                gts_arousal = gts[:, 1]

                rmse = torch.sqrt(tmf.mean_squared_error(preds, gts))
                r2 = tmf.r2_score(preds, gts)

                # Calculate metrics for valence
                rmse_valence = torch.sqrt(tmf.mean_squared_error(preds_valence, gts_valence))
                r2_valence = tmf.r2_score(preds_valence, gts_valence)

                # Calculate metrics for arousal
                rmse_arousal = torch.sqrt(tmf.mean_squared_error(preds_arousal, gts_arousal))
                r2_arousal = tmf.r2_score(preds_arousal, gts_arousal)

                log.info(f'*** Display RMSE and R² scores ({dataset.upper()}) ***')
                log.info(f"RMSE: {round(rmse.item(), 4)}")
                log.info(f"R²: {round(r2.item(), 4)}")
                log.info(f"Valence - RMSE: {round(rmse_valence.item(), 4)}, R²: {round(r2_valence.item(), 4)}")
                log.info(f"Arousal - RMSE: {round(rmse_arousal.item(), 4)}, R²: {round(r2_arousal.item(), 4)}")

                if self.output_file is not None:
                    with open(self.output_file, 'a') as f:
                        f.write(f"RMSE ({dataset.upper()}): {round(rmse.item(), 4)}\n")
                        f.write(f"R² ({dataset.upper()}): {round(r2.item(), 4)}\n")
                        f.write(f"Valence - RMSE ({dataset.upper()}): {round(rmse_valence.item(), 4)}\n")
                        f.write(f"Valence - R² ({dataset.upper()}): {round(r2_valence.item(), 4)}\n")
                        f.write(f"Arousal - RMSE ({dataset.upper()}): {round(rmse_arousal.item(), 4)}\n")
                        f.write(f"Arousal - R² ({dataset.upper()}): {round(r2_arousal.item(), 4)}\n")

                output_dic[f"test_rmse_{dataset}"] = round(rmse.item(), 4)
                output_dic[f"test_r2_{dataset}"] = round(r2.item(), 4)
                output_dic[f"test_rmse_valence_{dataset}"] = round(rmse_valence.item(), 4)
                output_dic[f"test_r2_valence_{dataset}"] = round(r2_valence.item(), 4)
                output_dic[f"test_rmse_arousal_{dataset}"] = round(rmse_arousal.item(), 4)
                output_dic[f"test_r2_arousal_{dataset}"] = round(r2_arousal.item(), 4)

        # Clear results for each dataset
        for dataset in ["jamendo",  "deam", "emomusic", "pmemo"]:
            if hasattr(self, f'{dataset}_results'):
                getattr(self, f'{dataset}_results')['preds'].clear()
                getattr(self, f'{dataset}_results')['gt'].clear()
                getattr(self, f'{dataset}_results')['paths'].clear()

        return output_dic

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.lr)
    
    def get_auc(self, prd_array, gt_array):
        prd_array = np.array(prd_array)
        gt_array = np.array(gt_array)

        prd_tensor = torch.tensor(prd_array)
        gt_tensor = torch.tensor(gt_array)
        
        try:
            roc_auc = tmf.auroc(prd_tensor, gt_tensor.int(), task='multilabel', num_labels = 56 , average='macro', num_classes=gt_tensor.size(1))
            pr_auc = tmf.average_precision(prd_tensor, gt_tensor.int(), task='multilabel', num_labels = 56, average='macro', num_classes=gt_tensor.size(1))            
        except ValueError as e:
            print(f"Error computing metrics: {e}")
            roc_auc = None
            pr_auc = None
        return roc_auc, pr_auc