""" train2.py — Ablation study: 6 experiments (backbone comparison) Pipeline IDENTICAL to train.py (train on full train_set, validate on official val_set, save best model by lowest val loss, 20 epochs). Only the model architecture differs across experiments. All models use ImageNet pretrained weights + standard VA regression head + CCC loss. Experiments: 1) VGG-16 + CCC only 2) MobileNetV2 + CCC only 3) ResNet-18 + CCC only 4) EfficientNetV2-S + CCC only 5) Swin-T + CCC only 6) ConvNeXt-Tiny + CCC only Output per experiment: experiments// ├── model.pt ├── train_history.csv └── *.png (10 plots) Summary: experiments/ ├── all_experiments_summary.csv └── summary_*.png """ 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 (identical to train.py) # ============================================================================ SCRIPT_DIR = Path(__file__).resolve().parent 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 paths (identical to train.py) # ============================================================================ 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 (identical to train.py) # ============================================================================ 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("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) # ============================================================================ # 6 Experiment configurations (backbone comparison, all CCC loss) # ============================================================================ EXPERIMENTS = [ { "name": "exp1_vgg16_ccc", "desc": "VGG-16 + ImageNet + CCC only", "model_type": "vgg16", "loss_type": "ccc_only", }, { "name": "exp2_mobilenetv2_ccc", "desc": "MobileNetV2 + ImageNet + CCC only", "model_type": "mobilenetv2", "loss_type": "ccc_only", }, { "name": "exp3_resnet18_ccc", "desc": "ResNet-18 + ImageNet + CCC only", "model_type": "resnet18", "loss_type": "ccc_only", }, { "name": "exp4_efficientnetv2_ccc", "desc": "EfficientNetV2-S + ImageNet + CCC only", "model_type": "efficientnetv2", "loss_type": "ccc_only", }, { "name": "exp5_swin_ccc", "desc": "Swin-T + ImageNet + CCC only", "model_type": "swin", "loss_type": "ccc_only", }, { "name": "exp6_convnext_ccc", "desc": "ConvNeXt-Tiny + ImageNet + CCC only", "model_type": "convnext", "loss_type": "ccc_only", }, ] # ============================================================================ # Plot helpers (from train.py) # ============================================================================ 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() 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() # ============================================================================ # Dataset & transforms (identical to train.py) # ============================================================================ 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 builders — all ImageNet pretrained, standard VA regression heads # ============================================================================ def _make_head(in_features, dropout=0.3): """Standard VA regression head used by all models.""" return nn.Sequential( nn.LayerNorm(in_features), nn.Linear(in_features, in_features), nn.Tanh(), nn.Dropout(dropout), nn.Linear(in_features, 2, bias=False), ) def build_model(exp_config): """ Build model based on experiment config. All models use ImageNet pretrained weights + standard VA regression head. """ model_type = exp_config["model_type"] if model_type == "vgg16": # VGG-16 model = models.vgg16(weights="DEFAULT") in_feat = model.classifier[6].in_features # 4096 model.classifier[6] = _make_head(in_feat) elif model_type == "mobilenetv2": # MobileNetV2 model = models.mobilenet_v2(weights="DEFAULT") in_feat = model.classifier[1].in_features # 1280 model.classifier = _make_head(in_feat) elif model_type == "resnet18": # ResNet-18 model = models.resnet18(weights="DEFAULT") in_feat = model.fc.in_features # 512 model.fc = _make_head(in_feat) elif model_type == "efficientnetv2": # EfficientNetV2-S model = models.efficientnet_v2_s(weights="DEFAULT") in_feat = model.classifier[1].in_features # 1280 model.classifier = _make_head(in_feat) elif model_type == "swin": # Swin Transformer Tiny model = models.swin_t(weights="DEFAULT") in_feat = model.head.in_features # 768 model.head = _make_head(in_feat) elif model_type == "convnext": # ConvNeXt Tiny model = models.convnext_tiny(weights="DEFAULT") in_feat = model.classifier[2].in_features # 768 model.classifier[2] = _make_head(in_feat) else: raise ValueError(f"Unknown model_type: {model_type}") model.to(DEVICE) return model # ============================================================================ # Loss functions & metrics (from train.py) # ============================================================================ val_loss_fn = nn.MSELoss() aro_loss_fn = nn.MSELoss() def CCCLoss(x, y): """CCC-based loss (returns -CCC). Identical to train.py.""" 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): """Compute training/validation loss based on experiment type.""" if loss_type == "mse_ccc": return ( 3 * val_loss_fn(val_pred, val_true) + 3 * aro_loss_fn(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): """CCC metric (numpy). Identical to train.py.""" 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 (same plots as train.py) # ============================================================================ def save_experiment_plots(history_df, exp_name, exp_dir): """Save the same set of plots as train.py for each experiment.""" epochs = history_df["epoch"].tolist() # 1. Train Loss save_line_plot( epochs, [history_df["train_loss"].tolist()], ["Train Loss"], f"{exp_name} — Train Loss", "Epoch", "Loss", exp_dir / "train_loss.png", ) # 2. Validation Loss save_line_plot( epochs, [history_df["valid_loss"].tolist()], ["Validation Loss"], f"{exp_name} — Validation Loss", "Epoch", "Loss", exp_dir / "valid_loss.png", ) # 3. Train vs Validation Loss 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", ) # 4. CCC_v 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", ) # 5. CCC_a 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", ) # 6. CCC_mean 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", ) # 7. All CCC on one plot 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", ) # 8. Learning Rate save_line_plot( epochs, [history_df["lr"].tolist()], ["Learning Rate"], f"{exp_name} — Learning Rate", "Epoch", "LR", exp_dir / "lr.png", ) # 9. Train Loss vs Valid Loss scatter save_scatter( history_df["train_loss"].tolist(), history_df["valid_loss"].tolist(), f"{exp_name} — Train Loss vs Valid Loss", "Train Loss", "Validation Loss", exp_dir / "train_vs_valid_scatter.png", annotate_labels=epochs, ) # 10. CCC_v vs CCC_a scatter save_scatter( history_df["ccc_v"].tolist(), history_df["ccc_a"].tolist(), f"{exp_name} — CCC_v vs CCC_a", "CCC_v", "CCC_a", exp_dir / "ccc_v_vs_ccc_a_scatter.png", annotate_labels=epochs, ) # ============================================================================ # Summary plots across all experiments # ============================================================================ def save_summary_plots(summary_df, save_dir): """Save comparison plots across all experiments.""" names = summary_df["name"].tolist() # Bar: Best CCC_v 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", ) # Bar: Best CCC_a 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", ) # Bar: Best CCC_mean 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", ) # Bar: Best Validation Loss save_bar_plot( names, summary_df["best_valid_loss"].tolist(), "Best Validation Loss by Experiment", "Experiment", "Loss", save_dir / "summary_best_valid_loss.png", ) # Bar: Best Epoch 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(): hist_csv = save_dir / row["name"] / "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: Validation Loss curves from all experiments plt.figure(figsize=(12, 6)) for _, row in summary_df.iterrows(): hist_csv = save_dir / row["name"] / "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() # Overlay: Train Loss curves from all experiments plt.figure(figsize=(12, 6)) for _, row in summary_df.iterrows(): hist_csv = save_dir / row["name"] / "train_history.csv" if hist_csv.exists(): h = pd.read_csv(hist_csv) plt.plot( h["epoch"], h["train_loss"], marker="o", linewidth=2, markersize=3, label=row["name"], ) plt.title("Train 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_train_loss_overlay.png", dpi=200) plt.close() # ============================================================================ # Training function — mirrors train.py's training loop EXACTLY # ============================================================================ def run_experiment(exp_config, train_loader, valid_loader): """ Run one experiment. Training loop structure is identical to train.py: - AdamW optimizer, CosineAnnealingLR(T_max=BATCHSIZE*NUM_EPOCHS) - AMP GradScaler + autocast - l2_reg/l1_reg computed per batch (not added to loss, kept for parity) - scheduler.step() called per batch (inside autocast, same as train.py) - Best model selected by minimum validation loss """ 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(exp_dir / "model.pt") print(f"\n{'='*70}") print(f"EXPERIMENT: {name}") print(f" {desc}") print(f" model_type={exp_config['model_type']}, " f"loss_type={loss_type}") print(f" output: {exp_dir}") print(f"{'='*70}") # Build model model = build_model(exp_config) # Optimizer, scheduler, scaler — identical to train.py 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 l2_lambda = 0.00001 # kept for parity with train.py (not added to loss) l1_lambda = 0.00001 history = { "epoch": [], "train_loss": [], "valid_loss": [], "ccc_v": [], "ccc_a": [], "ccc_mean": [], "lr": [], } for epoch in range(NUM_EPOCHS): # ── Train (identical structure to train.py) ────────────────────── 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}/{NUM_EPOCHS}" ): 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 = 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 (identical structure to train.py) ─────────────────── 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, 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 = 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_avg = 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_avg) 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"Training Loss: {train_loss:.4f}, " f"Validation Loss: {valid_loss_avg:.4f}, " f"CCC_v: {valid_ccc_v:.4f}, " f"CCC_a: {valid_ccc_a:.4f}, " f"CCC_mean: {valid_ccc_mean:.4f}, " f"Learning Rate: {current_lr:.8f}" ) if valid_loss_avg < best_valid_loss: best_valid_loss = valid_loss_avg 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 metrics 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, "model_type": exp_config["model_type"], "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)}") print(f"\n--- Running {len(EXPERIMENTS)} experiments ---") 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" + "=" * 70) print("ALL 6 EXPERIMENTS COMPLETE") print("=" * 70) print(summary_df.to_string(index=False)) print(f"\nSummary CSV: {EXPERIMENTS_DIR / 'all_experiments_summary.csv'}") print(f"All output: {EXPERIMENTS_DIR}")