""" train3.py – CAGE pipeline: 5-Fold StratifiedKFold (1 seed per fold) Pipeline: - train_set is split into 80% train / 20% val via StratifiedKFold - Official val_set is used as test set (evaluated after each fold) - Model: MaxViT-T + Standard head + Combined weights - Loss: 3*MSE_val + 3*MSE_aro + CCC_val + CCC_aro """ import pandas as pd import os import numpy as np import torch import torchvision.transforms as transforms import torchvision.models as models from torch.utils.data import DataLoader, Dataset import torch.nn as nn import torch.optim as optim from PIL import Image from torch.optim import lr_scheduler from tqdm import tqdm from pathlib import Path from sklearn.model_selection import StratifiedKFold import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt SCRIPT_DIR = Path(__file__).resolve().parent MODEL_PATH = os.getenv("MODEL_PATH", str(SCRIPT_DIR / "model.pt")) PLOTS_DIR = SCRIPT_DIR / "plots_cage" PLOTS_DIR.mkdir(parents=True, exist_ok=True) def resolve_images_dir(split_name, explicit_env_var, data_root, extract_root): explicit = os.getenv(explicit_env_var, "").strip() candidates = [] if explicit: candidates.append(Path(explicit)) candidates.append(Path(data_root) / f"{split_name}_set" / "images") candidates.append(Path(extract_root) / f"{split_name}_extracted" / "images") extract_split_root = Path(extract_root) / f"{split_name}_extracted" if extract_split_root.exists(): for p in extract_split_root.rglob("images"): if p.is_dir(): candidates.append(p) for p in candidates: if p.exists() and p.is_dir(): return str(p) tried = "\n".join([str(p) for p in candidates]) raise FileNotFoundError( f"Could not find images folder for split='{split_name}'. Tried:\n{tried}" ) # Load the annotations for training and validation from separate CSV files DATA_ROOT = os.getenv("AFFECTNET_ROOT", "/workspace/data_affectnet/AffectNet") EXTRACT_ROOT = os.getenv("AFFECTNET_EXTRACT_ROOT", f"{DATA_ROOT}/extracted") ANNO_ROOT = os.getenv("AFFECTNET_ANNO_ROOT", "../../affectnet_annotations") IMAGE_FOLDER = resolve_images_dir("train", "AFFECTNET_TRAIN_IMAGES", DATA_ROOT, EXTRACT_ROOT) IMAGE_FOLDER_TEST = resolve_images_dir("val", "AFFECTNET_VAL_IMAGES", DATA_ROOT, EXTRACT_ROOT) train_annotations_path = os.getenv( "AFFECTNET_TRAIN_ANNO", f"{ANNO_ROOT}/train_set_annotation_without_lnd.csv" ) valid_annotations_path = os.getenv( "AFFECTNET_VAL_ANNO", f"{ANNO_ROOT}/val_set_annotation_without_lnd.csv" ) train_annotations_df = pd.read_csv(train_annotations_path) valid_annotations_df = pd.read_csv(valid_annotations_path) # Set parameters BATCHSIZE = int(os.getenv("BATCHSIZE", "64")) NUM_EPOCHS = 20 LR = 4e-5 NUM_WORKERS = int(os.getenv("NUM_WORKERS", "0")) DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") print("DATA_ROOT =", DATA_ROOT) print("EXTRACT_ROOT =", EXTRACT_ROOT) print("AFFECTNET_TRAIN_IMAGES =", IMAGE_FOLDER) print("AFFECTNET_VAL_IMAGES =", IMAGE_FOLDER_TEST) print("MODEL_PATH =", MODEL_PATH) print("PLOTS_DIR =", PLOTS_DIR) print("Train set size =", len(train_annotations_df)) print("Official val_set size (final test only) =", len(valid_annotations_df)) # ========================= # Plot helpers # ========================= def save_line_plot(x, ys, labels, title, xlabel, ylabel, save_path): plt.figure(figsize=(10, 6)) for y, label in zip(ys, labels): plt.plot(x, y, marker="o", linewidth=2, markersize=4, label=label) plt.title(title) plt.xlabel(xlabel) plt.ylabel(ylabel) plt.grid(True, alpha=0.3) plt.legend() plt.tight_layout() plt.savefig(save_path, dpi=200) plt.close() def save_bar_plot(labels, values, title, xlabel, ylabel, save_path, annotate=True): plt.figure(figsize=(10, 6)) bars = plt.bar(labels, values) plt.title(title) plt.xlabel(xlabel) plt.ylabel(ylabel) plt.grid(True, axis="y", alpha=0.3) if annotate: for bar, val in zip(bars, values): plt.text( bar.get_x() + bar.get_width() / 2, bar.get_height(), f"{val:.4f}", ha="center", va="bottom", fontsize=9, ) plt.tight_layout() plt.savefig(save_path, dpi=200) plt.close() def save_boxplot(data_list, labels, title, ylabel, save_path): plt.figure(figsize=(10, 6)) plt.boxplot(data_list, labels=labels, vert=True) plt.title(title) plt.ylabel(ylabel) plt.grid(True, axis="y", alpha=0.3) plt.tight_layout() plt.savefig(save_path, dpi=200) plt.close() def save_histogram(values, bins, title, xlabel, ylabel, save_path): plt.figure(figsize=(10, 6)) plt.hist(values, bins=bins) plt.title(title) plt.xlabel(xlabel) plt.ylabel(ylabel) plt.grid(True, alpha=0.3) plt.tight_layout() plt.savefig(save_path, dpi=200) plt.close() def save_scatter(x, y, title, xlabel, ylabel, save_path, annotate_labels=None): plt.figure(figsize=(10, 6)) plt.scatter(x, y) if annotate_labels is not None: for xi, yi, txt in zip(x, y, annotate_labels): plt.annotate(str(txt), (xi, yi), fontsize=9) plt.title(title) plt.xlabel(xlabel) plt.ylabel(ylabel) plt.grid(True, alpha=0.3) plt.tight_layout() plt.savefig(save_path, dpi=200) plt.close() def save_fold_plots(history_df, fold_idx, seed, fold_dir): epochs = history_df["epoch"].tolist() save_line_plot( epochs, [history_df["train_loss"].tolist()], ["Train Loss"], f"CAGE Fold {fold_idx} Seed {seed} - Train Loss", "Epoch", "Loss", fold_dir / f"fold_{fold_idx}_train_loss.png", ) save_line_plot( epochs, [history_df["valid_loss"].tolist()], ["Validation Loss"], f"CAGE Fold {fold_idx} Seed {seed} - Validation Loss", "Epoch", "Loss", fold_dir / f"fold_{fold_idx}_valid_loss.png", ) save_line_plot( epochs, [history_df["train_loss"].tolist(), history_df["valid_loss"].tolist()], ["Train Loss", "Validation Loss"], f"CAGE Fold {fold_idx} Seed {seed} - Train vs Validation Loss", "Epoch", "Loss", fold_dir / f"fold_{fold_idx}_train_valid_loss.png", ) save_line_plot( epochs, [history_df["ccc_v"].tolist()], ["CCC_v"], f"CAGE Fold {fold_idx} Seed {seed} - CCC_v", "Epoch", "CCC", fold_dir / f"fold_{fold_idx}_ccc_v.png", ) save_line_plot( epochs, [history_df["ccc_a"].tolist()], ["CCC_a"], f"CAGE Fold {fold_idx} Seed {seed} - CCC_a", "Epoch", "CCC", fold_dir / f"fold_{fold_idx}_ccc_a.png", ) save_line_plot( epochs, [history_df["ccc_mean"].tolist()], ["CCC_mean"], f"CAGE Fold {fold_idx} Seed {seed} - CCC_mean", "Epoch", "CCC", fold_dir / f"fold_{fold_idx}_ccc_mean.png", ) save_line_plot( epochs, [ history_df["ccc_v"].tolist(), history_df["ccc_a"].tolist(), history_df["ccc_mean"].tolist(), ], ["CCC_v", "CCC_a", "CCC_mean"], f"CAGE Fold {fold_idx} Seed {seed} - Combined CCC Curves", "Epoch", "CCC", fold_dir / f"fold_{fold_idx}_ccc_all.png", ) save_line_plot( epochs, [history_df["lr"].tolist()], ["Learning Rate"], f"CAGE Fold {fold_idx} Seed {seed} - Learning Rate", "Epoch", "LR", fold_dir / f"fold_{fold_idx}_lr.png", ) save_scatter( history_df["train_loss"].tolist(), history_df["valid_loss"].tolist(), f"CAGE Fold {fold_idx} Seed {seed} - Train Loss vs Valid Loss", "Train Loss", "Validation Loss", fold_dir / f"fold_{fold_idx}_train_vs_valid_scatter.png", annotate_labels=epochs, ) save_scatter( history_df["ccc_v"].tolist(), history_df["ccc_a"].tolist(), f"CAGE Fold {fold_idx} Seed {seed} - CCC_v vs CCC_a", "CCC_v", "CCC_a", fold_dir / f"fold_{fold_idx}_ccc_v_vs_ccc_a.png", annotate_labels=epochs, ) def save_summary_plots(summary_df, save_dir): fold_labels = [f"Fold {x}" for x in summary_df["fold"].tolist()] save_bar_plot( fold_labels, summary_df["test_loss"].tolist(), "CAGE Final Test Loss by Fold", "Fold", "Loss", save_dir / "summary_test_loss_by_fold.png", ) save_bar_plot( fold_labels, summary_df["test_ccc_v"].tolist(), "CAGE Final Test CCC_v by Fold", "Fold", "CCC_v", save_dir / "summary_test_ccc_v_by_fold.png", ) save_bar_plot( fold_labels, summary_df["test_ccc_a"].tolist(), "CAGE Final Test CCC_a by Fold", "Fold", "CCC_a", save_dir / "summary_test_ccc_a_by_fold.png", ) save_bar_plot( fold_labels, summary_df["test_ccc_mean"].tolist(), "CAGE Final Test CCC_mean by Fold", "Fold", "CCC_mean", save_dir / "summary_test_ccc_mean_by_fold.png", ) save_bar_plot( fold_labels, summary_df["best_epoch"].tolist(), "CAGE Best Epoch by Fold", "Fold", "Best Epoch", save_dir / "summary_best_epoch_by_fold.png", annotate=True, ) save_boxplot( [ summary_df["test_ccc_v"].tolist(), summary_df["test_ccc_a"].tolist(), summary_df["test_ccc_mean"].tolist(), ], ["CCC_v", "CCC_a", "CCC_mean"], "CAGE Distribution of Final CCC Across 5 Folds", "CCC", save_dir / "summary_ccc_boxplot.png", ) save_histogram( summary_df["test_ccc_mean"].tolist(), bins=min(5, len(summary_df)), title="CAGE Histogram of Final CCC_mean", xlabel="CCC_mean", ylabel="Frequency", save_path=save_dir / "summary_ccc_mean_histogram.png", ) save_scatter( summary_df["best_epoch"].tolist(), summary_df["test_ccc_mean"].tolist(), "CAGE Best Epoch vs Final CCC_mean", "Best Epoch", "Final CCC_mean", save_dir / "summary_best_epoch_vs_ccc_mean.png", annotate_labels=summary_df["fold"].tolist(), ) # **** Create dataset and data loaders **** class CustomDataset(Dataset): def __init__(self, dataframe, root_dir, transform=None, balance=False): self.dataframe = dataframe self.transform = transform self.root_dir = root_dir self.balance = balance if self.balance: self.dataframe = self.balance_dataset() def __len__(self): return len(self.dataframe) def __getitem__(self, idx): img_id = int(float(self.dataframe["number"].iloc[idx])) image_path = os.path.join(self.root_dir, f"{img_id}.jpg") image = Image.open(image_path) classes = torch.tensor(self.dataframe.iloc[idx, 1], dtype=torch.int8) valence = torch.tensor(self.dataframe.iloc[idx, 2], dtype=torch.float16) arousal = torch.tensor(self.dataframe.iloc[idx, 3], dtype=torch.float16) if self.transform: image = self.transform(image) return image, classes, valence, arousal def balance_dataset(self): balanced_df = self.dataframe.groupby("exp", group_keys=False).apply( lambda x: x.sample(self.dataframe["exp"].value_counts().min()) ) return balanced_df transform = transforms.Compose( [ transforms.RandomHorizontalFlip(0.5), transforms.RandomGrayscale(0.01), transforms.RandomRotation(10), transforms.ColorJitter( brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1 ), transforms.RandomPerspective( distortion_scale=0.2, p=0.5 ), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), transforms.RandomErasing( p=0.5, scale=(0.02, 0.2), ratio=(0.3, 3.3), value="random" ), ] ) transform_valid = transforms.Compose( [ transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ] ) test_dataset = CustomDataset( dataframe=valid_annotations_df, root_dir=IMAGE_FOLDER_TEST, transform=transform_valid, balance=False, ) test_loader = DataLoader( test_dataset, batch_size=BATCHSIZE, shuffle=False, num_workers=NUM_WORKERS ) # ***** Define the model ***** def build_model(): MODEL = models.maxvit_t(weights="DEFAULT") block_channels = MODEL.classifier[3].in_features MODEL.classifier = nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.LayerNorm(block_channels), nn.Linear(block_channels, block_channels), nn.Tanh(), nn.Linear(block_channels, 10, bias=False), ) MODEL.to(DEVICE) COMBINED_MODEL_PATH = os.getenv( "COMBINED_MODEL_PATH", str((SCRIPT_DIR / "../AffectNet8_Maxvit_Combined/model.pt").resolve()), ) MODEL.load_state_dict(torch.load(COMBINED_MODEL_PATH, map_location=DEVICE)) MODEL.classifier = nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.LayerNorm(block_channels), nn.Linear(block_channels, block_channels), nn.Tanh(), nn.Dropout(0.3), nn.Linear(block_channels, 2, bias=False), ) MODEL.to(DEVICE) return MODEL def CCCLoss(x, y): x_mean = torch.mean(x, dim=0) y_mean = torch.mean(y, dim=0) x_var = torch.var(x, dim=0) y_var = torch.var(y, dim=0) cov_matrix = torch.matmul( (x - x_mean).permute(*torch.arange(x.dim() - 1, -1, -1)), y - y_mean ) / (x.size(0) - 1) numerator = 2 * cov_matrix denominator = x_var + y_var + torch.pow((x_mean - y_mean), 2) ccc = torch.mean(numerator / denominator) return -ccc def ccc_numpy_1d(x, y, eps=1e-8): x = np.asarray(x).reshape(-1) y = np.asarray(y).reshape(-1) x_mean = x.mean() y_mean = y.mean() x_var = x.var() y_var = y.var() cov = ((x - x_mean) * (y - y_mean)).mean() return (2.0 * cov) / (x_var + y_var + (x_mean - y_mean) ** 2 + eps) def evaluate_loader(model, loader): model.eval() eval_loss = 0.0 eval_val_true = [] eval_val_pred = [] eval_aro_true = [] eval_aro_pred = [] with torch.no_grad(): for images, _, val_true, aro_true in loader: images, val_true, aro_true = ( images.to(DEVICE), val_true.to(DEVICE), aro_true.to(DEVICE), ) with torch.autocast(device_type="cuda", dtype=torch.float16): outputs = model(images) val_pred = outputs[:, 0] aro_pred = outputs[:, 1] loss = ( 3 * val_loss(val_pred.cuda(), val_true.cuda()) + 3 * aro_loss(aro_pred.cuda(), aro_true.cuda()) + CCCLoss(val_pred.cuda(), val_true.cuda()) + CCCLoss(aro_pred.cuda(), aro_true.cuda()) ) eval_loss += loss.item() eval_val_true.extend(val_true.detach().cpu().float().numpy()) eval_val_pred.extend(val_pred.detach().cpu().float().numpy()) eval_aro_true.extend(aro_true.detach().cpu().float().numpy()) eval_aro_pred.extend(aro_pred.detach().cpu().float().numpy()) ccc_v = ccc_numpy_1d(eval_val_pred, eval_val_true) ccc_a = ccc_numpy_1d(eval_aro_pred, eval_aro_true) ccc_mean = (ccc_v + ccc_a) / 2.0 return eval_loss / max(1, len(loader)), ccc_v, ccc_a, ccc_mean val_loss = nn.MSELoss() aro_loss = nn.MSELoss() # ***** Train the model ***** print("--- Start CAGE Stratified 5-Fold training (1 fold = 1 seed) ---") seeds = [11, 32, 57, 23, 42] fold_test_metrics = [] global_best_valid_loss = float("inf") global_best_model_path = MODEL_PATH l2_lambda = 0.00001 l1_lambda = 0.00001 for fold_idx, seed in enumerate(seeds, start=1): print(f"--- Fold {fold_idx}/5 with seed {seed} ---") fold_dir = PLOTS_DIR / f"fold_{fold_idx}_seed_{seed}" fold_dir.mkdir(parents=True, exist_ok=True) skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=seed) split_list = list(skf.split(train_annotations_df, train_annotations_df["exp"])) train_idx, val_idx = split_list[fold_idx - 1] fold_train_df = train_annotations_df.iloc[train_idx].reset_index(drop=True) fold_valid_df = train_annotations_df.iloc[val_idx].reset_index(drop=True) print(f"Fold train size = {len(fold_train_df)}, fold val size = {len(fold_valid_df)}") fold_train_df.to_csv(fold_dir / f"fold_{fold_idx}_train_split.csv", index=False) fold_valid_df.to_csv(fold_dir / f"fold_{fold_idx}_valid_split.csv", index=False) train_dataset = CustomDataset( dataframe=fold_train_df, root_dir=IMAGE_FOLDER, transform=transform, balance=True, ) valid_dataset = CustomDataset( dataframe=fold_valid_df, root_dir=IMAGE_FOLDER, transform=transform_valid, balance=False, ) train_loader = DataLoader( train_dataset, batch_size=BATCHSIZE, shuffle=True, num_workers=NUM_WORKERS ) valid_loader = DataLoader( valid_dataset, batch_size=BATCHSIZE, shuffle=False, num_workers=NUM_WORKERS ) MODEL = build_model() optimizer = optim.AdamW(MODEL.parameters(), lr=LR) cosine_scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=BATCHSIZE * NUM_EPOCHS) scaler = torch.cuda.amp.GradScaler() fold_best_valid_loss = float("inf") fold_best_epoch = -1 fold_model_path = str(SCRIPT_DIR / f"model_s{seed}_fold{fold_idx}.pt") history = { "epoch": [], "train_loss": [], "valid_loss": [], "ccc_v": [], "ccc_a": [], "ccc_mean": [], "lr": [], } for epoch in range(NUM_EPOCHS): MODEL.train() train_loss = 0.0 current_lr = optimizer.param_groups[0]["lr"] for images, _, val_true, aro_true in tqdm( train_loader, desc=f"CAGE Fold {fold_idx} train progress" ): images, val_true, aro_true = ( images.to(DEVICE), val_true.to(DEVICE), aro_true.to(DEVICE), ) optimizer.zero_grad() l2_reg = 0 l1_reg = 0 with torch.autocast(device_type="cuda", dtype=torch.float16): outputs = MODEL(images) val_pred = outputs[:, 0] aro_pred = outputs[:, 1] for param in MODEL.parameters(): l2_reg += torch.norm(param, 2) l1_reg += torch.norm(param, 1) loss = ( 3 * val_loss(val_pred.cuda(), val_true.cuda()) + 3 * aro_loss(aro_pred.cuda(), aro_true.cuda()) + CCCLoss(val_pred.cuda(), val_true.cuda()) + CCCLoss(aro_pred.cuda(), aro_true.cuda()) ) train_loss += loss.item() scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() cosine_scheduler.step() train_loss_mean = train_loss / len(train_loader) valid_loss_mean, ccc_v, ccc_a, ccc_mean = evaluate_loader(MODEL, valid_loader) history["epoch"].append(epoch + 1) history["train_loss"].append(train_loss_mean) history["valid_loss"].append(valid_loss_mean) history["ccc_v"].append(ccc_v) history["ccc_a"].append(ccc_a) history["ccc_mean"].append(ccc_mean) history["lr"].append(current_lr) print( f"Fold [{fold_idx}/5] Epoch [{epoch+1}/{NUM_EPOCHS}] - " f"Training Loss: {train_loss_mean:.4f}, " f"Validation Loss (fold 20%): {valid_loss_mean:.4f}, " f"CCC_v: {ccc_v:.4f}, " f"CCC_a: {ccc_a:.4f}, " f"CCC_mean: {ccc_mean:.4f}, " f"Learning Rate: {current_lr:.8f}, " ) if valid_loss_mean < fold_best_valid_loss: fold_best_valid_loss = valid_loss_mean fold_best_epoch = epoch + 1 torch.save(MODEL.state_dict(), fold_model_path) print(f"Saved best fold model: {fold_model_path}") if fold_best_valid_loss < global_best_valid_loss: global_best_valid_loss = fold_best_valid_loss global_best_model_path = fold_model_path history_df = pd.DataFrame(history) history_df.to_csv(fold_dir / f"fold_{fold_idx}_epoch_history.csv", index=False) save_fold_plots(history_df, fold_idx, seed, fold_dir) MODEL.load_state_dict(torch.load(fold_model_path, map_location=DEVICE)) test_loss_mean, test_ccc_v, test_ccc_a, test_ccc_mean = evaluate_loader(MODEL, test_loader) fold_test_metrics.append( ( seed, fold_idx, fold_best_epoch, fold_best_valid_loss, test_loss_mean, test_ccc_v, test_ccc_a, test_ccc_mean, ) ) print( f"Fold [{fold_idx}/5] Final Test (official val_set) - " f"Loss: {test_loss_mean:.4f}, " f"CCC_v: {test_ccc_v:.4f}, " f"CCC_a: {test_ccc_a:.4f}, " f"CCC_mean: {test_ccc_mean:.4f}" ) summary_df = pd.DataFrame( fold_test_metrics, columns=[ "seed", "fold", "best_epoch", "best_valid_loss", "test_loss", "test_ccc_v", "test_ccc_a", "test_ccc_mean", ], ) summary_df.to_csv(PLOTS_DIR / "fold_summary.csv", index=False) save_summary_plots(summary_df, PLOTS_DIR) print("--- Final summary over runs on official val_set ---") test_losses = summary_df["test_loss"].to_numpy() test_ccc_vs = summary_df["test_ccc_v"].to_numpy() test_ccc_as = summary_df["test_ccc_a"].to_numpy() test_ccc_means = summary_df["test_ccc_mean"].to_numpy() print( f"Runs: {len(fold_test_metrics)}; Mean Loss: {np.mean(test_losses):.4f} (std {np.std(test_losses):.4f}), " f"Mean CCC_v: {np.mean(test_ccc_vs):.4f} (std {np.std(test_ccc_vs):.4f}), " f"Mean CCC_a: {np.mean(test_ccc_as):.4f} (std {np.std(test_ccc_as):.4f}), " f"Mean CCC_mean: {np.mean(test_ccc_means):.4f} (std {np.std(test_ccc_means):.4f})" ) torch.save(torch.load(global_best_model_path, map_location=DEVICE), MODEL_PATH) print(f"Best fold model copied to MODEL_PATH: {MODEL_PATH}") print(f"All plots and CSV logs saved to: {PLOTS_DIR}")