train.py

import argparse
import json
from pathlib import Path
from time import time

import torch
from torch import nn, optim

from dataset import create_dataloaders
from models import VGG16CIFAR, ResNet18CIFAR


def accuracy(logits, targets):
    preds = torch.argmax(logits, dim=1)
    correct = (preds == targets).sum().item()
    return correct / targets.size(0)


def train_one_epoch(model, loader, criterion, optimizer, device):
    model.train()
    total_loss = 0.0
    total_acc = 0.0
    total_samples = 0

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

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        batch_size = labels.size(0)
        total_loss += loss.item() * batch_size
        total_acc += accuracy(outputs, labels) * batch_size
        total_samples += batch_size

    return total_loss / total_samples, total_acc / total_samples


@torch.no_grad()
def evaluate(model, loader, criterion, device):
    model.eval()
    total_loss = 0.0
    total_acc = 0.0
    total_samples = 0

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

        outputs = model(images)
        loss = criterion(outputs, labels)

        batch_size = labels.size(0)
        total_loss += loss.item() * batch_size
        total_acc += accuracy(outputs, labels) * batch_size
        total_samples += batch_size

    return total_loss / total_samples, total_acc / total_samples


def build_model(name: str, num_classes: int):
    if name.lower() == "vgg16":
        return VGG16CIFAR(num_classes=num_classes)
    if name.lower() == "resnet18":
        return ResNet18CIFAR(num_classes=num_classes)
    raise ValueError(f"Unknown model: {name}")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=str, required=True, choices=["vgg16", "resnet18"])
    parser.add_argument("--epochs", type=int, default=10)
    parser.add_argument("--data-dir", type=str, default="Multi-class Weather Dataset")
    parser.add_argument("--batch-size", type=int, default=64)
    parser.add_argument("--lr", type=float, default=0.001)
    parser.add_argument("--val-split", type=float, default=0.2)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--num-workers", type=int, default=0)
    parser.add_argument("--output-dir", type=str, default="outputs")
    args = parser.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    train_loader, val_loader, class_to_idx = create_dataloaders(
        base_dir=".",
        data_dir=args.data_dir,
        batch_size=args.batch_size,
        val_split=args.val_split,
        num_workers=args.num_workers,
        seed=args.seed,
    )

    num_classes = len(class_to_idx)
    model = build_model(args.model, num_classes=num_classes).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=args.lr)

    history = {
        "model": args.model,
        "num_classes": num_classes,
        "train_loss": [],
        "train_acc": [],
        "val_loss": [],
        "val_acc": [],
        "epoch_time_sec": [],
    }

    for epoch in range(1, args.epochs + 1):
        start = time()
        train_loss, train_acc = train_one_epoch(
            model, train_loader, criterion, optimizer, device
        )
        val_loss, val_acc = evaluate(model, val_loader, criterion, device)
        elapsed = time() - start

        history["train_loss"].append(train_loss)
        history["train_acc"].append(train_acc)
        history["val_loss"].append(val_loss)
        history["val_acc"].append(val_acc)
        history["epoch_time_sec"].append(elapsed)

        print(
            f"Epoch {epoch:03d} | "
            f"train loss {train_loss:.4f} acc {train_acc:.4f} | "
            f"val loss {val_loss:.4f} acc {val_acc:.4f} | "
            f"time {elapsed:.1f}s"
        )

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    model_path = output_dir / f"{args.model}.pt"
    history_path = output_dir / f"{args.model}_history.json"
    class_path = output_dir / "classes.json"

    torch.save(model.state_dict(), model_path)
    history_path.write_text(json.dumps(history, indent=2), encoding="utf-8")
    class_path.write_text(json.dumps(class_to_idx, indent=2), encoding="utf-8")

    print(f"Saved model to {model_path}")
    print(f"Saved history to {history_path}")
    print(f"Saved classes to {class_path}")


if __name__ == "__main__":
    main()