""" train2.py – Ablation study: 6 experiments (3 models × 2 losses) Models: 1) MaxViT-T + AttentivePoolingHead + load Combined weights (attn_combined) 2) MaxViT-T + AttentivePoolingHead + NO Combined weights (attn_nocombined) 3) MaxViT-T + Standard head + NO Combined weights (std_nocombined) Losses: A) 3*MSE_val + 3*MSE_aro + CCC_val + CCC_aro (mse_ccc) B) CCC_val + CCC_aro (ccc_only) Same pipeline as train.py for each experiment. Each saves its own folder with per-epoch CSV, plots, and model checkpoint. """ 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 import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt # ========================= # Paths # ========================= SCRIPT_DIR = Path(__file__).resolve().parent MODEL_PATH = os.getenv("MODEL_PATH", str(SCRIPT_DIR / "model.pt")) EXPERIMENTS_DIR = SCRIPT_DIR / "experiments" EXPERIMENTS_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}" ) # ========================= # Data # ========================= 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) # ========================= # 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") COMBINED_MODEL_PATH = os.getenv( "COMBINED_MODEL_PATH", str((SCRIPT_DIR / "../AffectNet8_Maxvit_Combined/model.pt").resolve()), ) 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("COMBINED_MODEL_PATH=", COMBINED_MODEL_PATH) print("EXPERIMENTS_DIR =", EXPERIMENTS_DIR) print("Train set size =", len(train_annotations_df)) print("Val set size =", len(valid_annotations_df)) print("DEVICE =", DEVICE) print("BATCH_SIZE =", BATCHSIZE) print("NUM_EPOCHS =", NUM_EPOCHS) print("LR =", LR) # ========================= # Experiment configurations # ========================= EXPERIMENTS = [ { "name": "exp1_attn_combined_mse_ccc", "desc": "MaxViT-T + Attention + Combined weights + MSE+CCC", "use_attention": True, "load_combined": True, "loss_type": "mse_ccc", }, { "name": "exp2_attn_nocombined_mse_ccc", "desc": "MaxViT-T + Attention + ImageNet only + MSE+CCC", "use_attention": True, "load_combined": False, "loss_type": "mse_ccc", }, { "name": "exp3_standard_nocombined_mse_ccc", "desc": "MaxViT-T + Standard head + ImageNet only + MSE+CCC", "use_attention": False, "load_combined": False, "loss_type": "mse_ccc", }, { "name": "exp4_attn_combined_ccc_only", "desc": "MaxViT-T + Attention + Combined weights + CCC only", "use_attention": True, "load_combined": True, "loss_type": "ccc_only", }, { "name": "exp5_attn_nocombined_ccc_only", "desc": "MaxViT-T + Attention + ImageNet only + CCC only", "use_attention": True, "load_combined": False, "loss_type": "ccc_only", }, { "name": "exp6_standard_nocombined_ccc_only", "desc": "MaxViT-T + Standard head + ImageNet only + CCC only", "use_attention": False, "load_combined": False, "loss_type": "ccc_only", }, ] # ========================= # 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=(12, 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=8, ) plt.xticks(rotation=30, ha="right") plt.tight_layout() plt.savefig(save_path, dpi=200) plt.close() # ========================= # Dataset & transforms # ========================= 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]), ] ) # ========================= # Model components # ========================= class AttentivePoolingHead(nn.Module): def __init__(self, feat_dim, hidden_dim, dropout=0.3): super().__init__() self.attn_proj = nn.Linear(feat_dim, 1) self.mlp = nn.Sequential( nn.LayerNorm(feat_dim), nn.Dropout(dropout), nn.Linear(feat_dim, hidden_dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(hidden_dim, 2, bias=False), ) def forward(self, x): tokens = x.flatten(2).transpose(1, 2) # [B, HW, C] weights = torch.softmax(self.attn_proj(tokens), dim=1) # [B, HW, 1] pooled = (tokens * weights).sum(dim=1) # [B, C] return torch.tanh(self.mlp(pooled)) def build_model(use_attention, load_combined): model = models.maxvit_t(weights="DEFAULT") block_channels = model.classifier[3].in_features if load_combined: # Temporarily set classifier to match Combined model architecture 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) model.load_state_dict(torch.load(COMBINED_MODEL_PATH, map_location=DEVICE)) # Now set the final head if use_attention: model.classifier = AttentivePoolingHead( feat_dim=block_channels, hidden_dim=block_channels, dropout=0.3, ) else: 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 # ========================= # Loss & metrics # ========================= mse_criterion_val = nn.MSELoss() mse_criterion_aro = nn.MSELoss() 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 compute_loss(val_pred, aro_pred, val_true, aro_true, loss_type): if loss_type == "mse_ccc": return ( 3 * mse_criterion_val(val_pred, val_true) + 3 * mse_criterion_aro(aro_pred, aro_true) + CCCLoss(val_pred, val_true) + CCCLoss(aro_pred, aro_true) ) else: # ccc_only return CCCLoss(val_pred, val_true) + CCCLoss(aro_pred, aro_true) 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) # ========================= # Per-experiment plots # ========================= def save_experiment_plots(history_df, exp_name, exp_dir): epochs = history_df["epoch"].tolist() save_line_plot( epochs, [history_df["train_loss"].tolist()], ["Train Loss"], f"{exp_name} - Train Loss", "Epoch", "Loss", exp_dir / "train_loss.png", ) save_line_plot( epochs, [history_df["valid_loss"].tolist()], ["Validation Loss"], f"{exp_name} - Validation Loss", "Epoch", "Loss", exp_dir / "valid_loss.png", ) save_line_plot( epochs, [history_df["train_loss"].tolist(), history_df["valid_loss"].tolist()], ["Train Loss", "Validation Loss"], f"{exp_name} - Train vs Validation Loss", "Epoch", "Loss", exp_dir / "train_vs_valid_loss.png", ) save_line_plot( epochs, [history_df["ccc_v"].tolist()], ["CCC_v"], f"{exp_name} - Validation CCC_v", "Epoch", "CCC", exp_dir / "valid_ccc_v.png", ) save_line_plot( epochs, [history_df["ccc_a"].tolist()], ["CCC_a"], f"{exp_name} - Validation CCC_a", "Epoch", "CCC", exp_dir / "valid_ccc_a.png", ) save_line_plot( epochs, [history_df["ccc_mean"].tolist()], ["CCC_mean"], f"{exp_name} - Validation CCC_mean", "Epoch", "CCC", exp_dir / "valid_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"{exp_name} - All CCC Curves", "Epoch", "CCC", exp_dir / "valid_ccc_all.png", ) save_line_plot( epochs, [history_df["lr"].tolist()], ["Learning Rate"], f"{exp_name} - Learning Rate", "Epoch", "LR", exp_dir / "lr.png", ) # ========================= # Summary plots across experiments # ========================= def save_summary_plots(summary_df, save_dir): names = summary_df["name"].tolist() save_bar_plot( names, summary_df["best_ccc_v"].tolist(), "Best CCC_v by Experiment", "Experiment", "CCC_v", save_dir / "summary_best_ccc_v.png", ) save_bar_plot( names, summary_df["best_ccc_a"].tolist(), "Best CCC_a by Experiment", "Experiment", "CCC_a", save_dir / "summary_best_ccc_a.png", ) save_bar_plot( names, summary_df["best_ccc_mean"].tolist(), "Best CCC_mean by Experiment", "Experiment", "CCC_mean", save_dir / "summary_best_ccc_mean.png", ) save_bar_plot( names, summary_df["best_valid_loss"].tolist(), "Best Validation Loss by Experiment", "Experiment", "Loss", save_dir / "summary_best_valid_loss.png", ) save_bar_plot( names, summary_df["best_epoch"].tolist(), "Best Epoch by Experiment", "Experiment", "Epoch", save_dir / "summary_best_epoch.png", ) # Overlay CCC_mean curves from all experiments plt.figure(figsize=(12, 6)) for _, row in summary_df.iterrows(): exp_dir = save_dir / row["name"] hist_csv = exp_dir / "train_history.csv" if hist_csv.exists(): h = pd.read_csv(hist_csv) plt.plot( h["epoch"], h["ccc_mean"], marker="o", linewidth=2, markersize=3, label=row["name"], ) plt.title("Validation CCC_mean - All Experiments") plt.xlabel("Epoch") plt.ylabel("CCC_mean") plt.grid(True, alpha=0.3) plt.legend(fontsize=8) plt.tight_layout() plt.savefig(save_dir / "summary_ccc_mean_overlay.png", dpi=200) plt.close() # Overlay valid loss curves from all experiments plt.figure(figsize=(12, 6)) for _, row in summary_df.iterrows(): exp_dir = save_dir / row["name"] hist_csv = exp_dir / "train_history.csv" if hist_csv.exists(): h = pd.read_csv(hist_csv) plt.plot( h["epoch"], h["valid_loss"], marker="o", linewidth=2, markersize=3, label=row["name"], ) plt.title("Validation Loss - All Experiments") plt.xlabel("Epoch") plt.ylabel("Loss") plt.grid(True, alpha=0.3) plt.legend(fontsize=8) plt.tight_layout() plt.savefig(save_dir / "summary_valid_loss_overlay.png", dpi=200) plt.close() # ========================= # Training function # ========================= def run_experiment(exp_config, train_loader, valid_loader): name = exp_config["name"] desc = exp_config["desc"] loss_type = exp_config["loss_type"] exp_dir = EXPERIMENTS_DIR / name exp_dir.mkdir(parents=True, exist_ok=True) exp_model_path = str(SCRIPT_DIR / f"model_{name}.pt") print(f"\n{'='*60}") print(f"EXPERIMENT: {name}") print(f" {desc}") print(f" use_attention={exp_config['use_attention']}, " f"load_combined={exp_config['load_combined']}, " f"loss_type={loss_type}") print(f" model_path={exp_model_path}") print(f" plots_dir={exp_dir}") print(f"{'='*60}") model = build_model(exp_config["use_attention"], exp_config["load_combined"]) optimizer = optim.AdamW(model.parameters(), lr=LR) scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=BATCHSIZE * NUM_EPOCHS) scaler = torch.cuda.amp.GradScaler() best_valid_loss = float("inf") best_epoch = -1 history = { "epoch": [], "train_loss": [], "valid_loss": [], "ccc_v": [], "ccc_a": [], "ccc_mean": [], "lr": [], } for epoch in range(NUM_EPOCHS): # ── Train ────────────────────────────────────────────── 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"[{name}] Epoch {epoch+1} train" ): images = images.to(DEVICE) val_true = val_true.to(DEVICE) aro_true = aro_true.to(DEVICE) optimizer.zero_grad() with torch.autocast(device_type="cuda", dtype=torch.float16): outputs = model(images) val_pred = outputs[:, 0] aro_pred = outputs[:, 1] loss = compute_loss(val_pred, aro_pred, val_true, aro_true, loss_type) train_loss += loss.item() scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() scheduler.step() train_loss /= max(1, len(train_loader)) # ── Validate ─────────────────────────────────────────── model.eval() valid_loss_sum = 0.0 valid_val_true = [] valid_val_pred = [] valid_aro_true = [] valid_aro_pred = [] with torch.no_grad(): for images, _, val_true, aro_true in valid_loader: images = images.to(DEVICE) val_true = val_true.to(DEVICE) aro_true = 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 = compute_loss( val_pred, aro_pred, val_true, aro_true, loss_type ) valid_loss_sum += loss.item() valid_val_true.extend(val_true.detach().cpu().float().numpy()) valid_val_pred.extend(val_pred.detach().cpu().float().numpy()) valid_aro_true.extend(aro_true.detach().cpu().float().numpy()) valid_aro_pred.extend(aro_pred.detach().cpu().float().numpy()) valid_loss = valid_loss_sum / max(1, len(valid_loader)) valid_ccc_v = ccc_numpy_1d(valid_val_pred, valid_val_true) valid_ccc_a = ccc_numpy_1d(valid_aro_pred, valid_aro_true) valid_ccc_mean = (valid_ccc_v + valid_ccc_a) / 2.0 history["epoch"].append(epoch + 1) history["train_loss"].append(train_loss) history["valid_loss"].append(valid_loss) history["ccc_v"].append(valid_ccc_v) history["ccc_a"].append(valid_ccc_a) history["ccc_mean"].append(valid_ccc_mean) history["lr"].append(current_lr) print( f"[{name}] Epoch [{epoch+1}/{NUM_EPOCHS}] - " f"Train Loss: {train_loss:.4f}, " f"Valid Loss: {valid_loss:.4f}, " f"CCC_v: {valid_ccc_v:.4f}, " f"CCC_a: {valid_ccc_a:.4f}, " f"CCC_mean: {valid_ccc_mean:.4f}, " f"LR: {current_lr:.8f}" ) if valid_loss < best_valid_loss: best_valid_loss = valid_loss best_epoch = epoch + 1 print(f" -> Saving best model at epoch {best_epoch}") torch.save(model.state_dict(), exp_model_path) # ── Save CSV & plots ─────────────────────────────────────── history_df = pd.DataFrame(history) history_df.to_csv(exp_dir / "train_history.csv", index=False) save_experiment_plots(history_df, name, exp_dir) # Print best epoch info best_idx = best_epoch - 1 print(f"\n=== [{name}] Best Epoch Metrics ===") print(f"Best epoch: {best_epoch}") print(f"Train Loss: {history_df['train_loss'][best_idx]:.4f}") print(f"Valid Loss: {history_df['valid_loss'][best_idx]:.4f}") print(f"Valid CCC_v: {history_df['ccc_v'][best_idx]:.4f}") print(f"Valid CCC_a: {history_df['ccc_a'][best_idx]:.4f}") print(f"Valid CCC_mean: {history_df['ccc_mean'][best_idx]:.4f}") return { "name": name, "desc": desc, "use_attention": exp_config["use_attention"], "load_combined": exp_config["load_combined"], "loss_type": exp_config["loss_type"], "best_epoch": best_epoch, "best_valid_loss": float(history_df["valid_loss"][best_idx]), "best_ccc_v": float(history_df["ccc_v"][best_idx]), "best_ccc_a": float(history_df["ccc_a"][best_idx]), "best_ccc_mean": float(history_df["ccc_mean"][best_idx]), } # ========================= # Main # ========================= print("\n--- Creating datasets (shared across all experiments) ---") train_dataset = CustomDataset( dataframe=train_annotations_df, root_dir=IMAGE_FOLDER, transform=transform, balance=True, ) valid_dataset = CustomDataset( dataframe=valid_annotations_df, root_dir=IMAGE_FOLDER_TEST, 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 ) print(f"Train dataset size (balanced): {len(train_dataset)}") print(f"Valid dataset size: {len(valid_dataset)}") all_results = [] for exp_config in EXPERIMENTS: result = run_experiment(exp_config, train_loader, valid_loader) all_results.append(result) # ========================= # Summary # ========================= summary_df = pd.DataFrame(all_results) summary_df.to_csv(EXPERIMENTS_DIR / "all_experiments_summary.csv", index=False) save_summary_plots(summary_df, EXPERIMENTS_DIR) print("\n" + "=" * 60) print("ALL EXPERIMENTS COMPLETE") print("=" * 60) print(summary_df.to_string(index=False)) print(f"\nSummary CSV: {EXPERIMENTS_DIR / 'all_experiments_summary.csv'}") print(f"All plots: {EXPERIMENTS_DIR}")