train.py

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import wandb

from data.load_data import CropData 
from models.cnn.resnet import get_resnet18
from models.vit.deit import get_deit3
from engine.losses import FocalLoss
from engine.trainer import train_model, get_optimal_device

import argparse
import yaml

def main():
    parser = argparse.ArgumentParser(description="Training Script for COL780 Assignment 3")
    parser.add_argument("--config", type=str, required=True, help="Path to YAML config file")
    args = parser.parse_args()

    with open(args.config, 'r') as f:
        config = yaml.safe_load(f)

    wandb.init(project="COL7680-A3", name=f"{config.get('backbone', 'Model')}_{config.get('loss_function', 'Loss')}", config=config)
    device = get_optimal_device()

    img_dir = "data/A3_Dataset"
    train_dataset = CropData(img_dir_path=img_dir, csv_file_path=f"{img_dir}/train.csv")
    val_dataset = CropData(img_dir_path=img_dir, csv_file_path=f"{img_dir}/validation.csv")
    test_dataset = CropData(img_dir_path=img_dir, csv_file_path=f"{img_dir}/test.csv")
    
    batch_size = config.get("batch_size", 32)
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=0)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)

    num_classes = config.get("num_classes", 10)
    if config.get("model_family", "cnn").lower() == "cnn":
        model = get_resnet18(num_classes=num_classes, use_se=config.get("use_se", False))
    else:
        model = get_deit3(num_classes=num_classes, use_dyt=config.get("use_dyt", False), init_alpha=config.get("init_alpha", 0.1))

    if config.get("loss_function") == "Focal":
        criterion = FocalLoss(gamma=config.get("focal_gamma", 2.0))
    else:
        criterion = nn.CrossEntropyLoss()

    model.to(device)
    
    lr = config.get("learning_rate", 1e-4)
    wd = config.get("weight_decay", 1e-4)
    if config.get("optimizer") == "AdamW":
        optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=wd)
    else:
        optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=wd)

    scheduler_type = config.get("scheduler", "None")
    epochs = config.get("epochs", 10)
    scheduler = None
    if scheduler_type == "OneCycleLR":
        scheduler = optim.lr_scheduler.OneCycleLR(
            optimizer, max_lr=lr * 10, 
            steps_per_epoch=len(train_loader), epochs=epochs
        )
    elif scheduler_type == "CosineAnnealingLR":
        warmup_epochs = 3 # 3-4 warmup epochs
        warmup_scheduler = optim.lr_scheduler.LinearLR(optimizer, start_factor=0.01, end_factor=1.0, total_iters=warmup_epochs)
        cosine_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs - warmup_epochs)
        scheduler = optim.lr_scheduler.SequentialLR(optimizer, schedulers=[warmup_scheduler, cosine_scheduler], milestones=[warmup_epochs])

    train_model(
        model=model, 
        train_loader=train_loader, 
        val_loader=val_loader, 
        criterion=criterion, 
        optimizer=optimizer, 
        scheduler=scheduler,
        epochs=epochs, 
        device=device
    )

    wandb.finish()

if __name__ == "__main__":
    main()