train.py

import os
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, WeightedRandomSampler
from sklearn.model_selection import StratifiedKFold, train_test_split
from sklearn.metrics import f1_score
from tqdm import tqdm

from config import (DEVICE, SAVE_DIR, TRAIN_DIR, BATCH_SIZE,
                    LR, EPOCHS, K_FOLDS, NUM_WORKERS, NUM_GROUPS)
from dataset import WBCDataset, get_transforms


class FocalLoss(nn.Module):
    def __init__(self, alpha=1, gamma=2):
        super().__init__()
        self.alpha, self.gamma = alpha, gamma
        self.ce = nn.CrossEntropyLoss(label_smoothing=0.1)

    def forward(self, inputs, targets):
        ce_loss = self.ce(inputs, targets)
        pt = torch.exp(-ce_loss)
        return self.alpha * (1 - pt) ** self.gamma * ce_loss


def train_convnext_kfold(df, le, model_class, use_clahe=True, tag="model"):
    os.makedirs(SAVE_DIR, exist_ok=True)
    train_tf, val_tf = get_transforms()
    skf = StratifiedKFold(n_splits=K_FOLDS, shuffle=True, random_state=42)

    for fold, (train_idx, val_idx) in enumerate(skf.split(df, df['label_encoded'])):
        print(f"\nFold {fold+1}/{K_FOLDS}")
        train_df, val_df = df.iloc[train_idx], df.iloc[val_idx]

        w = 1. / np.bincount(train_df['label_encoded'])
        sampler = WeightedRandomSampler(w[train_df['label_encoded']], len(train_df))

        train_loader = DataLoader(
            WBCDataset(train_df, TRAIN_DIR, train_tf, use_clahe=use_clahe),
            batch_size=BATCH_SIZE, sampler=sampler, num_workers=NUM_WORKERS)
        val_loader = DataLoader(
            WBCDataset(val_df, TRAIN_DIR, val_tf, use_clahe=use_clahe),
            batch_size=BATCH_SIZE, num_workers=NUM_WORKERS)

        model = model_class(NUM_GROUPS, len(le.classes_)).to(DEVICE)
        optimizer = optim.AdamW(model.parameters(), lr=LR, weight_decay=0.05)
        scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS)
        criterion = FocalLoss()

        best_f1 = 0
        for epoch in range(EPOCHS):
            model.train()
            for imgs, groups, fines in tqdm(train_loader, desc=f"Ep {epoch+1}"):
                imgs   = imgs.to(DEVICE)
                groups = groups.to(DEVICE)
                fines  = fines.to(DEVICE)
                optimizer.zero_grad()
                out_g, out_f = model(imgs)
                loss = criterion(out_g, groups) + criterion(out_f, fines)
                loss.backward()
                optimizer.step()
            scheduler.step()

            model.eval()
            all_p, all_t = [], []
            with torch.no_grad():
                for imgs, _, fines in val_loader:
                    _, out_f = model(imgs.to(DEVICE))
                    all_p.extend(out_f.argmax(1).cpu().numpy())
                    all_t.extend(fines.numpy())

            score = f1_score(all_t, all_p, average='macro')
            print(f"Fold {fold+1} Ep {epoch+1} | F1: {score:.4f}")

            if score > best_f1:
                best_f1 = score
                path = f"{SAVE_DIR}/best_{tag}_fold{fold+1}.pth"
                torch.save(model.state_dict(), path)



def train_simple(df, le, model_class, tag="resnet"):
    """Entrainement simple (sans K-Fold) """
    

    os.makedirs(SAVE_DIR, exist_ok=True)
    train_tf, val_tf = get_transforms()

    train_df, val_df = train_test_split(
        df, test_size=0.15, stratify=df['label_encoded'], random_state=42)

    counts  = np.bincount(train_df['label_encoded'])
    weights = torch.tensor(1.0 / np.log1p(counts), dtype=torch.float).to(DEVICE)
    weights = weights / weights.sum() * len(le.classes_)

    train_loader = DataLoader(
        WBCDataset(train_df, TRAIN_DIR, train_tf),
        batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS)
    val_loader = DataLoader(
        WBCDataset(val_df, TRAIN_DIR, val_tf),
        batch_size=BATCH_SIZE, num_workers=NUM_WORKERS)

    model     = model_class(len(le.classes_)).to(DEVICE)
    optimizer = optim.AdamW(model.parameters(), lr=LR, weight_decay=0.05)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS)
    criterion = nn.CrossEntropyLoss(weight=weights, label_smoothing=0.1)

    best_f1 = 0
    for epoch in range(EPOCHS):
        model.train()
        for imgs, _, labels in tqdm(train_loader, desc=f"Ep {epoch+1}"):
            imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)
            optimizer.zero_grad()
            loss = criterion(model(imgs), labels)
            loss.backward()
            optimizer.step()
        scheduler.step()

        model.eval()
        all_p, all_t = [], []
        with torch.no_grad():
            for imgs, _, labels in val_loader:
                out = model(imgs.to(DEVICE))
                all_p.extend(out.argmax(1).cpu().numpy())
                all_t.extend(labels.numpy())

        score = f1_score(all_t, all_p, average='macro')
        print(f"Ep {epoch+1} | F1: {score:.4f}")

        if score > best_f1:
            best_f1 = score
            torch.save(
                {'model': model.state_dict(), 'le': le},
                f"{SAVE_DIR}/best_{tag}.pth")
            
def finetune_convnext_kfold(df, le, model_class, src_tag="convnext", dst_tag="convnext_ft",
                             lr=2e-5, epochs=10, use_clahe=True):
    """
    Repart des meilleurs poids K-Fold puis lr plus petit pour affiner les classes difficiles.
    """
    os.makedirs(SAVE_DIR, exist_ok=True)
    _, val_tf = get_transforms()
    train_tf, _ = get_transforms()
    skf = StratifiedKFold(n_splits=K_FOLDS, shuffle=True, random_state=42)

    for fold, (train_idx, val_idx) in enumerate(skf.split(df, df['label_encoded'])):
        src_path = f"{SAVE_DIR}/best_{src_tag}_fold{fold+1}.pth"
        if not os.path.exists(src_path):
            print(f"Fold {fold+1} introuvable ({src_path}), ignoré.")
            continue

        print(f"\nFine-tuning Fold {fold+1}/{K_FOLDS}")
        train_df, val_df = df.iloc[train_idx], df.iloc[val_idx]

        w = 1. / np.bincount(train_df['label_encoded'])
        sampler = WeightedRandomSampler(w[train_df['label_encoded']], len(train_df))

        train_loader = DataLoader(
            WBCDataset(train_df, TRAIN_DIR, train_tf, use_clahe=use_clahe),
            batch_size=BATCH_SIZE, sampler=sampler, num_workers=NUM_WORKERS)
        val_loader = DataLoader(
            WBCDataset(val_df, TRAIN_DIR, val_tf, use_clahe=use_clahe),
            batch_size=BATCH_SIZE, num_workers=NUM_WORKERS)

        # Charge les poids existants
        model = model_class(NUM_GROUPS, len(le.classes_)).to(DEVICE)
        model.load_state_dict(torch.load(src_path, map_location=DEVICE))

        optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=0.05)
        scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
        criterion = FocalLoss()

        best_f1 = 0
        for epoch in range(epochs):
            model.train()
            for imgs, groups, fines in tqdm(train_loader, desc=f"FT Fold {fold+1} Ep {epoch+1}"):
                imgs   = imgs.to(DEVICE)
                groups = groups.to(DEVICE)
                fines  = fines.to(DEVICE)
                optimizer.zero_grad()
                out_g, out_f = model(imgs)
                loss = criterion(out_g, groups) + criterion(out_f, fines)
                loss.backward()
                optimizer.step()
            scheduler.step()

            model.eval()
            all_p, all_t = [], []
            with torch.no_grad():
                for imgs, _, fines in val_loader:
                    _, out_f = model(imgs.to(DEVICE))
                    all_p.extend(out_f.argmax(1).cpu().numpy())
                    all_t.extend(fines.numpy())

            score = f1_score(all_t, all_p, average='macro')
            print(f"FT Fold {fold+1} Ep {epoch+1} | F1: {score:.4f}")

            if score > best_f1:
                best_f1 = score
                torch.save(model.state_dict(), f"{SAVE_DIR}/best_{dst_tag}_fold{fold+1}.pth")