# import pandas as pd # import os # 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 # SCRIPT_DIR = Path(__file__).resolve().parent # MODEL_PATH = os.getenv("MODEL_PATH", str(SCRIPT_DIR / "model.pt")) # 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) # # **** 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): # image_path = os.path.join( # self.root_dir, f"{self.dataframe['number'].iloc[idx]}.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 # ), # model more robust to changes in lighting conditions. # transforms.RandomPerspective( # distortion_scale=0.2, p=0.5 # ), # can be helpful if your images might have varying perspectives. # transforms.ToTensor(), # saves image as tensor (automatically divides by 255) # 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" # ), # Should help overfitting # ] # ) # transform_valid = transforms.Compose( # [ # transforms.ToTensor(), # transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), # ] # ) # 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 # ) # # ***** Define the model ***** # # Initialize the 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) # def CCCLoss(x, y): # # Compute means # x_mean = torch.mean(x, dim=0) # y_mean = torch.mean(y, dim=0) # # Compute variances # x_var = torch.var(x, dim=0) # y_var = torch.var(y, dim=0) # # Compute covariance matrix # cov_matrix = torch.matmul( # (x - x_mean).permute(*torch.arange(x.dim() - 1, -1, -1)), y - y_mean # ) / (x.size(0) - 1) # # Compute CCC # numerator = 2 * cov_matrix # denominator = x_var + y_var + torch.pow((x_mean - y_mean), 2) # ccc = torch.mean(numerator / denominator) # return -ccc # val_loss = nn.MSELoss() # aro_loss = nn.MSELoss() # optimizer = optim.AdamW(MODEL.parameters(), lr=LR) # lr_scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=BATCHSIZE * NUM_EPOCHS) # # ***** Train the model ***** # print("--- Start training ---") # scaler = torch.cuda.amp.GradScaler() # best_valid_loss = float("inf") # l2_lambda = 0.00001 # L1 Regularization # l1_lambda = 0.00001 # L2 Regularization # for epoch in range(NUM_EPOCHS): # MODEL.train() # total_train_correct = 0 # total_train_samples = 0 # current_lr = optimizer.param_groups[0]["lr"] # for images, _, val_true, aro_true in tqdm( # train_loader, desc="Epoch train_loader progress" # ): # images, val_true, aro_true = ( # images.to(DEVICE), # val_true.to(DEVICE), # aro_true.to(DEVICE), # ) # optimizer.zero_grad() # train_loss = 0 # 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) # **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()) # ) # # + l2_lambda * l2_reg + l1_lambda * l1_reg # train_loss += loss.item() # scaler.scale(loss).backward() # scaler.step(optimizer) # scaler.update() # lr_scheduler.step() # MODEL.eval() # valid_loss = 0.0 # total_valid_correct = 0 # total_valid_samples = 0 # 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 = ( # 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()) # ) # valid_loss += loss.item() # print( # f"Epoch [{epoch+1}/{NUM_EPOCHS}] - " # f"Training Loss: {train_loss/len(train_loader):.4f}, " # f"Validation Loss: {valid_loss/len(valid_loader):.4f}, " # f"Learning Rate: {current_lr:.8f}, " # ) # valid_loss_mean = valid_loss / max(1, len(valid_loader)) # if valid_loss_mean < best_valid_loss: # best_valid_loss = valid_loss_mean # print(f"Saving model at epoch {epoch+1}") # torch.save(MODEL.state_dict(), MODEL_PATH) # Save the best model # ── APViT repo path setup (must precede mmcls/mmcv imports) ───────────────── import sys from pathlib import Path as _Path _apvit_path = str(_Path(__file__).resolve().parent.parent.parent / "APViT") if _apvit_path not in sys.path: sys.path.insert(0, _apvit_path) # ───────────────────────────────────────────────────────────────────────────── import pandas as pd import os import numpy as np import torch import torchvision.transforms as transforms 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 mmcv from mmcls.models import build_classifier 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" PLOTS_DIR.mkdir(parents=True, exist_ok=True) APVIT_REPO = SCRIPT_DIR.parent.parent / "APViT" APVIT_CONFIG = APVIT_REPO / "configs" / "apvit" / "RAF.py" 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)) print("StratifiedKFold splits = 5 (each fold train/val = 8:2)") # ========================= # 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"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"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"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"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"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"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"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"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"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"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(), "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(), "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(), "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(), "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(), "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"], "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="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(), "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.Resize(112), # APViT / IR-50 requires 112x112 input 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.Resize(112), # APViT / IR-50 requires 112x112 input 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 ) APVIT_EMBED_DIM = 768 def _build_apvit_backbone(): """Build APViT PoolingVitClassifier (backbone only, no head).""" cfg = mmcv.Config.fromfile(str(APVIT_CONFIG)) cfg.model.pretrained = None cfg.model.head = None ir50_w = APVIT_REPO / "weights" / "backbone_ir50_ms1m_epoch120.pth" vit_small_w = APVIT_REPO / "weights" / "vit_small_p16_224-15ec54c9.pth" cfg.model.extractor.pretrained = str(ir50_w) if ir50_w.exists() else None cfg.model.vit.pretrained = str(vit_small_w) if vit_small_w.exists() else None if cfg.model.extractor.pretrained is None: print("[WARNING] IR-50 weights not found at", ir50_w, "— training from random init.") if cfg.model.vit.pretrained is None: print("[WARNING] ViT-Small weights not found at", vit_small_w, "— training from random init.") return build_classifier(cfg.model) class APViTWithHead(nn.Module): def __init__(self, num_outputs): super().__init__() self.apvit = _build_apvit_backbone() self.head = nn.Sequential( nn.LayerNorm(APVIT_EMBED_DIM), nn.Linear(APVIT_EMBED_DIM, APVIT_EMBED_DIM), nn.Tanh(), nn.Linear(APVIT_EMBED_DIM, num_outputs, bias=False), ) def forward(self, x): features, _ = self.apvit.extract_feat(x) # [B, 768] CLS token return self.head(features) COMBINED_MODEL_PATH = os.getenv( "COMBINED_MODEL_PATH", str((SCRIPT_DIR / "../AffectNet8_Maxvit_Combined/model.pt").resolve()), ) def build_model(): model = APViTWithHead(num_outputs=10) # Match Combined checkpoint structure model.to(DEVICE) model.load_state_dict(torch.load(COMBINED_MODEL_PATH, map_location=DEVICE)) # Replace head with VA regression head (2 outputs) model.head = nn.Sequential( nn.LayerNorm(APVIT_EMBED_DIM), nn.Linear(APVIT_EMBED_DIM, APVIT_EMBED_DIM), nn.Tanh(), nn.Dropout(0.3), nn.Linear(APVIT_EMBED_DIM, 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 training on full train_set, validate on val_set ---") l2_lambda = 0.00001 l1_lambda = 0.00001 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 ) model = build_model() 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): model.train() train_loss = 0.0 train_val_true = [] train_val_pred = [] train_aro_true = [] train_aro_pred = [] current_lr = optimizer.param_groups[0]["lr"] for images, _, val_true, aro_true in tqdm(train_loader, desc=f"Epoch {epoch+1} train_loader 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, val_true) + 3 * aro_loss(aro_pred, aro_true) + CCCLoss(val_pred, val_true) + CCCLoss(aro_pred, aro_true) ) train_loss += loss.item() scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() scheduler.step() train_val_true.extend(val_true.detach().cpu().numpy()) train_val_pred.extend(val_pred.detach().cpu().numpy()) train_aro_true.extend(aro_true.detach().cpu().numpy()) train_aro_pred.extend(aro_pred.detach().cpu().numpy()) train_loss /= max(1, len(train_loader)) train_ccc_v = ccc_numpy_1d(train_val_pred, train_val_true) train_ccc_a = ccc_numpy_1d(train_aro_pred, train_aro_true) train_ccc_mean = (train_ccc_v + train_ccc_a) / 2.0 model.eval() valid_loss = 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 = ( 3 * val_loss(val_pred, val_true) + 3 * aro_loss(aro_pred, aro_true) + CCCLoss(val_pred, val_true) + CCCLoss(aro_pred, aro_true) ) valid_loss += loss.item() valid_val_true.extend(val_true.detach().cpu().numpy()) valid_val_pred.extend(val_pred.detach().cpu().numpy()) valid_aro_true.extend(aro_true.detach().cpu().numpy()) valid_aro_pred.extend(aro_pred.detach().cpu().numpy()) valid_loss /= 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"Epoch [{epoch+1}/{NUM_EPOCHS}] - " f"Train Loss: {train_loss:.4f}, " f"Valid Loss: {valid_loss:.4f}, " f"Valid CCC_v: {valid_ccc_v:.4f}, " f"Valid CCC_a: {valid_ccc_a:.4f}, " f"Valid CCC_mean: {valid_ccc_mean:.4f}, " f"Learning Rate: {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(), MODEL_PATH) # Save training history and plot if needed import pandas as pd history_df = pd.DataFrame(history) history_df.to_csv(PLOTS_DIR / "train_history.csv", index=False) # Plotting (reuse your plot helpers, but only plot what makes sense) save_line_plot( history_df["epoch"].tolist(), [history_df["train_loss"].tolist()], ["Train Loss"], "Train Loss", "Epoch", "Loss", PLOTS_DIR / "train_loss.png", ) save_line_plot( history_df["epoch"].tolist(), [history_df["valid_loss"].tolist()], ["Validation Loss"], "Validation Loss", "Epoch", "Loss", PLOTS_DIR / "valid_loss.png", ) save_line_plot( history_df["epoch"].tolist(), [history_df["train_loss"].tolist(), history_df["valid_loss"].tolist()], ["Train Loss", "Validation Loss"], "Train vs Validation Loss", "Epoch", "Loss", PLOTS_DIR / "train_vs_valid_loss.png", ) save_line_plot( history_df["epoch"].tolist(), [history_df["ccc_v"].tolist()], ["CCC_v"], "Validation CCC_v", "Epoch", "CCC", PLOTS_DIR / "valid_ccc_v.png", ) save_line_plot( history_df["epoch"].tolist(), [history_df["ccc_a"].tolist()], ["CCC_a"], "Validation CCC_a", "Epoch", "CCC", PLOTS_DIR / "valid_ccc_a.png", ) save_line_plot( history_df["epoch"].tolist(), [history_df["ccc_mean"].tolist()], ["CCC_mean"], "Validation CCC_mean", "Epoch", "CCC", PLOTS_DIR / "valid_ccc_mean.png", ) save_line_plot( history_df["epoch"].tolist(), [history_df["lr"].tolist()], ["Learning Rate"], "Learning Rate", "Epoch", "LR", PLOTS_DIR / "lr.png", ) # In thông tin của epoch tốt nhất best_idx = best_epoch - 1 print("=== 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}") n_epochs_run = len(history_df["ccc_v"]) avg_ccc_v = sum(history_df["ccc_v"]) / n_epochs_run avg_ccc_a = sum(history_df["ccc_a"]) / n_epochs_run avg_ccc_mean = sum(history_df["ccc_mean"]) / n_epochs_run avg_valid_loss = sum(history_df["valid_loss"]) / n_epochs_run print("\n=== Average Across All Epochs ===") print(f"Avg Valid Loss : {avg_valid_loss:.4f}") print(f"Avg CCC_v : {avg_ccc_v:.4f}") print(f"Avg CCC_a : {avg_ccc_a:.4f}") print(f"Avg CCC_mean : {avg_ccc_mean:.4f}") print("=================================") print(f"--- Training complete. Best epoch: {best_epoch}, Best val loss: {best_valid_loss:.4f} ---") print(f"Best model saved to: {MODEL_PATH}") print(f"All plots and CSV logs saved to: {PLOTS_DIR}")