train_classifier.py

import torch
import wandb
import hydra
from tqdm import tqdm
import matplotlib.pyplot as plt
import os
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

from models.multimodalAttention import MultiModalAttentionClassifier
from models.dinov2 import DinoV2Finetune

from torch.optim.lr_scheduler import ReduceLROnPlateau

from utils.sanity import show_images
import numpy as np
import pandas as pd


@hydra.main(config_path="configs", config_name="train")
def main (cfg):
    print ("START")
    model = train (cfg)
    #test_model(cfg, model)



def train(cfg, train_idx=None, val_idx=None):
    """
        Fonction globale qui entraîne le réseau. Les paramètres à fixer pour le modèle sont dans le 
        fichier config/train.yaml. 
        
        Ces données config sont accessibles via le paramètre cfg qui n'est
        pas à renseigner lors de l'appel de la fonction.
    """
    
    base_dir = os.path.dirname(os.path.abspath(__file__))
    csv_path = os.path.join(base_dir, "dataset/log1pviews_per_category.csv")
    categories_df = pd.read_csv(csv_path)
    logger = (
        wandb.init(
            project="challenge_CSC_43M04_EP",
            name=cfg.experiment_name,
        )
        if cfg.log
        else None
    )

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print (device)
    # On crée le modèle défini dans train.yaml sur hydra et le to(device) le balance 
    # sur le cpu s'il existet
    #model = hydra.utils.instantiate(cfg.model.instance).to(device)
    model = MultiModalAttentionClassifier(classification_dim=len(categories_df)).to(device)
    # On crée l'optimizer défini sur train.yaml
    optimizer = hydra.utils.instantiate(cfg.optim, params=model.parameters())
    # Idem et le datamodule permet globalement de charger les images et les fournir au modèle
    datamodule = hydra.utils.instantiate(cfg.datamodule, train_idx=train_idx, val_idx=val_idx)
    train_loader = datamodule.train_dataloader()
    val_loader = datamodule.val_dataloader()
    extreme_train_loader = datamodule.extreme_train_dataloader()
    
    # Permet de charger le modèle avec le meilleur validation loss en cas 
    # de remontée du val_loss
    val_loss_min = np.inf


    # Le scheduler permet de réduire le learning rate en même temps que la loss du validation set diminue
    # mode = 'min' car on veut que le learning rate diminue
    # factor = 0.3 -> le learning rate est diminué de ce facteur quand la condition est remplie
    # patience = 3 : nombre d'epoch sans amélioration avant que le learning rate ne soit réduit
    # Les paramètres sont empiriques. Ne pas hésiter à modifier
    # On retient dans une variable min_learning rate le learning rate final du scheduler
    # Cette variable min_learning_rate est très importante car elle conditionne la fin de
    # la convergence
    scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=cfg.factor_learning_rate, patience=cfg.patience_learning_rate, min_lr=cfg.min_learning_rate)

    # Envoie le sanity check a wandb pour le training set
    csv_path = os.path.join(base_dir, "assets/sanity/train_images")
    train_sanity = show_images(train_loader, name=csv_path)
    (
        logger.log({"sanity_checks/train_images": wandb.Image(train_sanity)})
        if logger is not None
        else None
    )
    # Envoie le sanity check a wandb pour le validation set
    csv_path = os.path.join(base_dir, "assets/sanity/val_images")
    val_sanity = show_images(val_loader, name=csv_path)
    logger.log(
        {"sanity_checks/val_images": wandb.Image(val_sanity)}
    )
    
    # Helper to plot and log target distributions
    def log_target_distribution(loader, name, logger):
        all_targets = []
        for batch in loader:
            targets = batch["target"].cpu().numpy().flatten()
            all_targets.append(targets)
        all_targets = np.concatenate(all_targets)
        plt.figure()
        plt.hist(all_targets, bins=30, alpha=0.7)
        plt.title(f"Target Distribution: {name}")
        plt.xlabel("Target")
        plt.ylabel("Count")
        plt.tight_layout()
        csv_path = os.path.join(base_dir, f"assets/sanity/{name}_target_dist.png")
        plt.savefig(csv_path)
        if logger is not None:
            logger.log({f"sanity_checks/{name}_target_dist": wandb.Image(plt.gcf())})
        plt.close()

    log_target_distribution(train_loader, "train", logger)
    log_target_distribution(val_loader, "val", logger)
    # Le max_epoch est juste une sécurité et ne devrait pas influer sur la fin de la convergence
    max_epochs = cfg.max_epochs
    epoch = 0

    ##################
    # Enregistrement #
    ##################

    
    if False:
    #Uniquement si on souhaite restaurer un modèle qui était en entrainement    
        checkpoint = torch.load('checkpoints/MIN_ATT&DAR_MULTIMODAL_2025-05-24_23-37-11.pt', weights_only=False)
        model.load_state_dict(checkpoint['model_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
        epoch = checkpoint['epoch'] + 1  # Reprend à l'epoch suivante

    
    # La je compte les élèments du train_loader pour mettre en place une balance pour équilibrer 
    # les catégories dans le calcul de la loss
    all_classes = []
    for batch in train_loader:
        # Assure-toi que batch["class_target"] est sur CPU et converti en numpy
        class_targets = batch["class_target"].detach().cpu().numpy().flatten()
        all_classes.extend(class_targets)

    all_classes = np.array(all_classes)
    unique, counts = np.unique(all_classes, return_counts=True)
    print("Class distribution in train_loader:")
    for u, c in zip(unique, counts):
        print(f"Class {u}: {c} samples")

    # Calcul des poids inverses à la fréquence
    class_weights = 1. / counts
    class_weights = class_weights / class_weights.sum()  # Normalisation (optionnel)
    class_weights = torch.tensor(class_weights, dtype=torch.float32, device=device)
    print("Class weights:", class_weights)


    loss_fn = torch.nn.CrossEntropyLoss(weight=class_weights)


    print ("*********")
    print ("Début training loop")
    print ("********")

    # On interrompt la boucle en fonction du learning rate et du max_epoch. 
    # (cf min_learning_rate) plus haut
    # Cf condition break à la fin.
    while True : 
        #################
        # Training loop #
        #################

    
        model.train()
        epoch_train_loss = 0
        # Compte le nombre d'images entraînées pour faire la moyenne pour le train_loss
        num_samples_train = 0
        # Là c'est juste la barre de progression pour la console
        pbar = tqdm(train_loader, desc=f"Epoch {epoch}", leave=False)
        for i, batch in enumerate(pbar):
            # On envoie les images dans le GPU (si dispo)
            batch["image"] = batch["image"].to(device)
            # Pareil pour les labels
            batch["target"] = batch["target"].to(device).squeeze()
            batch["channel"] = batch["channel"].to(device)
            batch["year"] = batch["year"].to(device)
            batch["http_count"] = batch["http_count"].to(device)
            batch["diese"] = batch["diese"].to(device)
            batch["nb_mots"] = batch["nb_mots"].to(device)
            batch["class_target"] = batch["class_target"].long().to(device)
            # Pass forward
            preds = model(batch).squeeze()
            
            loss = loss_fn(preds, batch["class_target"])

            # Log weights, biases, and gradients to wandb
            if logger is not None:
                for name, param in model.named_parameters():
                    if param.requires_grad:
                        logger.log({f"weights/{name}": wandb.Histogram(param.detach().cpu().numpy())})
                        if param.grad is not None:
                            logger.log({f"grads/{name}": wandb.Histogram(param.grad.detach().cpu().numpy())})
            
            # Partie éliminée pour gagner en vitesse 
            # Là on envoit les données a wandb pour qu'il les affiche
            (
                logger.log({"loss": loss.detach().cpu().numpy(),"year_avg": batch["year"].cpu().numpy().mean()})
                if logger is not None
                else None
            )
            # Classico
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            epoch_train_loss += loss.detach().cpu().numpy() * len(batch["image"])
            num_samples_train += len(batch["image"])
            # Affiche la progression dans la console
            pbar.set_postfix({"train/loss_step": loss.detach().cpu().numpy(), "learning_rate": optimizer.param_groups[0]["lr"]})
        epoch_train_loss /= num_samples_train
        # Pareil, on envoit a wandb
        (
            logger.log(
                {
                    "epoch": epoch,
                    "train/loss_epoch": epoch_train_loss,
                }
            )
            if logger is not None
            else None
        )

        ###################
        # Validation loop #
        ###################

        val_metrics = {}
        epoch_val_loss = 0
        num_samples_val = 0
        model.eval()
        all_targets = []
        all_preds = []
        all_losses = []
        all_train_losses = []
        all_val_losses = []
        for _, batch in enumerate(val_loader):
            batch["image"] = batch["image"].to(device)
            # Pareil pour les labels
            batch["target"] = batch["target"].to(device).squeeze()
            batch["channel"] = batch["channel"].to(device)
            batch["year"] = batch["year"].to(device)
            batch["http_count"] = batch["http_count"].to(device)
            batch["diese"] = batch["diese"].to(device)
            batch["nb_mots"] = batch["nb_mots"].to(device)
            batch["class_target"] = batch["class_target"].long().to(device)
            with torch.no_grad():
                preds = model(batch)
            # On récupère la class prédite
            num_pred = torch.argmax(preds, dim=1)  # For classification, get the predicted class
            num_pred = categories_df["avg_log1p_views"].values[num_pred.cpu().numpy()]  # Convert to tensor and reshape
            num_pred = torch.tensor(num_pred, device=device, dtype=torch.float32)  # Convert to tensor and move to device
            loss = torch.nn.functional.mse_loss(num_pred, batch["target"], reduction='none')  # shape: (batch_size,)
            
            all_val_losses.append(loss.detach().cpu().numpy().squeeze())

            
            
            train_loss = torch.nn.functional.cross_entropy(preds, batch["class_target"], reduction='none')  # shape: (batch_size,)
            # Collect for scatter plot
            all_targets.append(batch["target"].detach().cpu().numpy())
            all_preds.append(num_pred.detach().cpu().numpy())
            all_losses.append(loss.detach().cpu().numpy().squeeze())
            all_train_losses.append(train_loss.detach().cpu().numpy())
            print("preds shape:", preds.shape)
            print("num_pred shape:", num_pred.shape)
            print("batch['target'] shape:", batch["target"].shape)
            print("batch['class_target'] shape:", batch["class_target"].shape)
            print("loss shape:", loss.shape)
            print("train_loss shape:", train_loss.shape)
        
        # After validation loop, scatter predictions vs target
        all_targets = np.concatenate(all_targets)
        all_preds = np.concatenate(all_preds)
        all_losses = np.concatenate(all_losses)
        all_train_losses = np.concatenate(all_train_losses)
        epoch_val_loss = all_losses.mean()
        high_val_loss = all_losses[all_losses > 10].mean() if len(all_losses[all_losses > 10]) > 0 else 0
        low_val_loss = all_losses[all_losses < 10].mean() if len(all_losses[all_losses < 10]) > 0 else 0

        # Plot 1: Predictions vs Target
        plt.figure(figsize=(6, 5))
        plt.scatter(all_targets, all_preds, alpha=0.5,c=all_losses, cmap='viridis')
        plt.plot([all_targets.min(), all_targets.max()], [all_targets.min(), all_targets.max()], 'r--')
        plt.xlabel("Target")
        plt.ylabel("Prediction")
        plt.title("Predictions vs Target (Validation)")
        plt.tight_layout()
        #plt.savefig("assets/sanity/val_pred_vs_target.png")
        if logger is not None:
            logger.log({f"predictions/val_pred_vs_target": wandb.Image(plt.gcf()),"epoch": epoch})
        plt.close()

        # Plot 2: Loss vs Target
        plt.figure(figsize=(6, 5))
        sc = plt.scatter(all_targets, all_losses, alpha=0.5, c=all_preds, cmap='viridis')
        plt.xlabel("Target")
        plt.ylabel("MSE Loss")
        plt.title("Loss vs Target (Validation)")
        plt.colorbar(sc, label="Prediction")
        plt.tight_layout()
        #plt.savefig("assets/sanity/val_loss_vs_target.png")
        if logger is not None:
            logger.log({f"predictions/val_loss_vs_target": wandb.Image(plt.gcf()), "epoch": epoch})
        plt.close()

        #plot 3: Train Loss vs Target
        plt.figure(figsize=(6, 5))
        sc = plt.scatter(all_targets, all_train_losses, alpha=0.5, c=all_preds, cmap='viridis')
        plt.xlabel("Target")
        plt.ylabel("Train Loss")
        plt.title("Train Loss vs Target (Validation)")
        plt.colorbar(sc, label="Prediction")
        plt.tight_layout()
        #plt.savefig("assets/sanity/val_train_loss_vs_target.png")
        if logger is not None:
            logger.log({f"predictions/val_train_loss_vs_target": wandb.Image(plt.gcf()), "epoch": epoch})
        plt.close()
        

        #epoch_val_loss = all_val_losses.mean()


        # On envoie la loss au scheduler pour qu'il puisse influencer le learning rate
        scheduler.step(epoch_val_loss)
        # On récupère le learning rate effectif du modèle avant de l'envoyer à wandb
        current_lr = optimizer.param_groups[0]["lr"]
        print ("Epoch : " + str(epoch) + ", Learning rate : " + str(current_lr)+ ", Training Loss : " + str(epoch_train_loss) + ", Validation Loss : " + str(epoch_val_loss))
        # On envoie tout à wandb
        val_metrics["val/loss_epoch"] = epoch_val_loss
        val_metrics["learning_rate"] = current_lr
        #val_metrics["val/high_loss"] = high_val_loss
        #val_metrics["val/low_loss"] = low_val_loss
        (
            logger.log(
                {
                    "epoch": epoch,
                    **val_metrics,
                }
            )
            if logger is not None
            else None
        )

        ##############
        # Sauvegarde #
        ##############

        if (epoch == 0) :   
            # 1ere epoch : on n'a pas encore enregistré de modèle
            val_loss_min = epoch_val_loss
        
        # On sauvegarde à intervalles réguliers
        if (epoch % cfg.checkpoint_interval == 0) :
            # Si oui, on sauvegarde le modèle    
            checkpoint = {
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'scheduler_state_dict': scheduler.state_dict(),
                'val_loss': epoch_val_loss,
            }
            torch.save(checkpoint, cfg.checkpoint_path)
            print ("Modèle enregistré !")

        if (epoch_val_loss <= val_loss_min) :    
            val_loss_min = epoch_val_loss
            # Si oui, on sauvegarde le modèle    
            checkpoint = {
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'scheduler_state_dict': scheduler.state_dict(),
                'val_loss': epoch_val_loss,
            }
            torch.save(checkpoint, cfg.min_checkpoint_path)
            print ("Modèle optimal enregistré !")


        # Si le modèle est pire que le meilleur, on repart avec le précédent
        if (epoch != 0 and val_loss_min < epoch_val_loss * (1 - cfg.aberration_val_loss)):
            print ("Aberration de la val_loss")
            checkpoint = torch.load(cfg.min_checkpoint_path, weights_only=False)
            model.load_state_dict(checkpoint['model_state_dict'])
            optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
            for param_group in optimizer.param_groups:
                param_group['lr'] *= cfg.factor_learning_rate
        

        ################################
        # Conditions de sortie de loop #
        ################################
        epoch += 1

        if (current_lr <= cfg.min_learning_rate or epoch > max_epochs) :
            print ("***")
            print ("Fin de la convergence")
            print ("***")
            break

    print(
        f"""Epoch {epoch}: 
        Training metrics:
        - Train Loss: {epoch_train_loss:.4f},
        Validation metrics: 
        - Val Loss: {epoch_val_loss:.4f}"""
    )

    if cfg.log:
        logger.finish()

    checkpoint = {
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'scheduler_state_dict': scheduler.state_dict()
    }
    torch.save(checkpoint, cfg.checkpoint_path)
    print ("Modèle enregistré !")
    return (model)





if __name__ == "__main__":
    main ()