train.py

import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, random_split
import matplotlib.pyplot as plt

from utils import CIFAR10Dataset, DataAugmentation
from model import ResNet


CONFIG = {
    "n": 18,
    "batch_size": 256,
    "max_lr": 5e-3,
    "epochs": 150,
    "weight_decay": 0.01,
}

ARTIFACT_DIR = "artifacts"
os.makedirs(ARTIFACT_DIR, exist_ok=True)


def get_device():
    if torch.backends.mps.is_available():
        return "mps"
    if torch.cuda.is_available():
        return "cuda"
    return "cpu"


def train():
    device = get_device()
    print(f"Running on: {device}")

    # Dataset + split
    full_train_set = CIFAR10Dataset(train=True)
    train_subset, val_subset = random_split(full_train_set, [45000, 5000])

    train_loader = DataLoader(
        train_subset,
        batch_size=CONFIG["batch_size"],
        shuffle=True,
        num_workers=12,
        pin_memory=True,
    )

    val_loader = DataLoader(
        val_subset,
        batch_size=CONFIG["batch_size"],
        shuffle=False,
        num_workers=4,
        pin_memory=True,
    )

    augmenter = DataAugmentation().to(device)
    augmenter.calculate_stats(train_loader,device)

    # Model + optimization
    torch.set_float32_matmul_precision("high")
    model = ResNet(n=CONFIG["n"]).to(device)
    model = torch.compile(model)

    optimizer = optim.AdamW(
        model.parameters(),
        lr=CONFIG["max_lr"] / 25,
        weight_decay=CONFIG["weight_decay"],
    )

    scheduler = optim.lr_scheduler.OneCycleLR(
        optimizer,
        max_lr=CONFIG["max_lr"],
        steps_per_epoch=len(train_loader),
        epochs=CONFIG["epochs"],
    )

    criterion = nn.CrossEntropyLoss()

    history = {
        "train_loss": [],
        "val_acc": [],
        "lr": [],
    }

    best_acc = 0.0

    for epoch in range(CONFIG["epochs"]):
        model.train()
        running_loss = 0.0

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

            images = augmenter(images, training=True)

            optimizer.zero_grad(set_to_none=True)
            outputs = model(images)
            loss = criterion(outputs, labels)

            loss.backward()
            optimizer.step()
            scheduler.step()

            running_loss += loss.item()
            history["lr"].append(scheduler.get_last_lr()[0])

        avg_loss = running_loss / len(train_loader)
        val_acc = evaluate(model, val_loader, augmenter, device)

        history["train_loss"].append(avg_loss)
        history["val_acc"].append(val_acc)

        print(
            f"Epoch {epoch + 1}/{CONFIG['epochs']} | "
            f"Loss: {avg_loss:.4f} | Val Acc: {val_acc:.2f}%"
        )

        if val_acc > best_acc:
            best_acc = val_acc
            torch.save(model.state_dict(), "best_resnet.pth")

    plot_training_artifacts(history)
    save_config(best_acc)


def evaluate(model, loader, augmenter, device):
    model.eval()
    correct = 0
    total = 0

    with torch.no_grad():
        for images, labels in loader:
            images = images.to(device)
            labels = labels.to(device)

            images = augmenter(images, training=False)
            outputs = model(images)

            preds = outputs.argmax(dim=1)
            total += labels.size(0)
            correct += (preds == labels).sum().item()

    return 100.0 * correct / total


def plot_training_artifacts(history):
    plt.figure()
    plt.plot(history["train_loss"])
    plt.title("Training Loss")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.savefig(f"{ARTIFACT_DIR}/loss_curve.png")
    plt.close()

    plt.figure()
    plt.plot(history["val_acc"])
    plt.title("Validation Accuracy")
    plt.xlabel("Epoch")
    plt.ylabel("Accuracy (%)")
    plt.savefig(f"{ARTIFACT_DIR}/val_accuracy.png")
    plt.close()

    plt.figure()
    plt.plot(history["lr"])
    plt.title("OneCycleLR Schedule")
    plt.xlabel("Iteration")
    plt.ylabel("Learning Rate")
    plt.savefig(f"{ARTIFACT_DIR}/lr_schedule.png")
    plt.close()


def save_config(best_acc):
    with open("hyperparameters.txt", "w") as f:
        for k, v in CONFIG.items():
            f.write(f"{k}: {v}\n")
        f.write(f"Best Val Acc: {best_acc:.2f}%\n")


if __name__ == "__main__":
    train()