train.py

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import numpy as np
import pandas as pd

from model import LocalizationModel
from dataset import NodeLocalizationDataset, collate_localization_batch
from config import HEAD_CLASSES


def masked_multihead_loss(head_logits, head_targets, sample_head_names, criterion):
    batch_size = len(sample_head_names)
    device = next(iter(head_logits.values())).device
    total_loss = torch.zeros(batch_size, device=device)

    for head_name, logits in head_logits.items():
        mask = torch.tensor(
            [name == head_name for name in sample_head_names],
            device=device, dtype=torch.bool,
        )
        if not mask.any():
            continue
        targets = head_targets[head_name].to(device)
        per_sample_loss = criterion(logits, targets)
        total_loss = total_loss + per_sample_loss * mask.float()
    
    return total_loss.sum() / batch_size

def train_model():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")

    num_epochs = 20
    batch_size = 32

    full_df = pd.read_csv('data.csv')
    
    train_df = full_df[full_df["session"] == 1]
    val_df = full_df[full_df["session"] == 2]

    train_csv_path = 'data_train_split.csv'
    val_csv_path = 'data_val_split.csv'
    train_df.to_csv(train_csv_path, index=False)
    val_df.to_csv(val_csv_path, index=False)

    train_dataset = NodeLocalizationDataset(
        csv_file=train_csv_path, root_dir='images/',
        topology_map_path='topology_map.json',
        apply_color_jitter=True)

    val_dataset = NodeLocalizationDataset(
        csv_file=val_csv_path, root_dir='images/',
        topology_map_path='topology_map.json',
        apply_color_jitter=False)

    train_nodes = set(train_df["node_id"].unique())
    val_nodes = set(val_df["node_id"].unique())
    missing_from_train = val_nodes - train_nodes
    if missing_from_train:
        print(f"WARNING: {len(missing_from_train)} node(s) appear in session 2 "
              f"but not session 1 -- these can never be predicted correctly: "
              f"{sorted(missing_from_train)[:10]}")

    print(f"unique nodes -> session 1 (train): {len(train_nodes)}, "
          f"session 2 (val): {len(val_nodes)}")
    print(f"Train images: {len(train_dataset)} | Val images: {len(val_dataset)}")

    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True,
                              collate_fn=collate_localization_batch)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False,
                            collate_fn=collate_localization_batch)

    model = LocalizationModel().to(device)

    for param in model.vision.parameters():
        param.requires_grad = False

    criterion = nn.CrossEntropyLoss(reduction="none")

    optimizer = optim.Adam([
        {'params': model.vision.parameters(), 'lr': 1e-5},
        {'params': model.shared.parameters(), 'lr': 1e-4},
        {'params': model.heads.parameters(), 'lr': 1e-4},
    ])

    best_val_loss = float('inf')

    for epoch in range(num_epochs):
        model.train()
        train_loss = 0.0

        if epoch == 5:
            print("Epoch 5. Starting Vision Backbone Finetuning")

            for param in model.vision.parameters():
                param.requires_grad = True

        for images, head_targets, sample_head_names in train_loader:
            images = images.to(device)

            optimizer.zero_grad()
            head_logits = model(images)
            loss = masked_multihead_loss(head_logits, head_targets, sample_head_names, criterion)
            loss.backward()
            optimizer.step()

            train_loss += loss.item()

        avg_loss = train_loss / len(train_loader)

        model.eval()
        val_loss = 0.0
        correct = {name: 0 for name in HEAD_CLASSES}
        total = {name: 0 for name in HEAD_CLASSES}

        with torch.no_grad():
            for val_images, val_head_targets, val_sample_head_names in val_loader:
                val_images = val_images.to(device)

                head_logits = model(val_images)
                loss = masked_multihead_loss(head_logits, val_head_targets, val_sample_head_names, criterion)

                val_loss += loss.item()

                for head_name in HEAD_CLASSES:
                    mask = [name == head_name for name in val_sample_head_names]
                    if not any(mask):
                        continue
                    mask_idx = torch.tensor(mask, dtype=torch.bool)
                    preds = head_logits[head_name][mask_idx].argmax(dim=1).cpu()
                    targets = val_head_targets[head_name][mask_idx]
                    correct[head_name] += (preds == targets).sum().item()
                    total[head_name] += mask_idx.sum().item()

        avg_val_loss = val_loss / len(val_loader)
        acc_str = " | ".join(
            f"{name}: {correct[name]/total[name]:.3f}" if total[name] > 0 else f"{name}: n/a"
            for name in HEAD_CLASSES
        )

        print(f"[ Epoch {epoch+1} ] train loss: {avg_loss:.4f} | val loss: {avg_val_loss:.4f}")
        print(f"  val acc -> {acc_str}")
 
        if avg_val_loss < best_val_loss:
            best_val_loss = avg_val_loss
            torch.save(model.state_dict(), 'localization_mobilenetv2.pth')
            print("best model saved!")

    print("train finished!")


if __name__ == '__main__':
    train_model()