train.py

"""Training entry point for ImageClassifier.

Usage examples::

    # Full training run
    python train.py --epochs 20 --batch 64 --lr 1e-3 --seed 42

    # Quick smoke test (1-batch overfit check)
    python train.py --test
"""

import argparse
import os
import random
import sys
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

# Ensure the project root is on sys.path so `src.*` imports work
_project_root = Path(__file__).resolve().parent
if str(_project_root) not in sys.path:
    sys.path.insert(0, str(_project_root))

from src.ml.model import TinyCNN
from src.ml.data import get_synthetic_dataset


# ---------------------------------------------------------------------------
# Reproducibility
# ---------------------------------------------------------------------------

def seed_everything(seed: int) -> None:
    """Set all random seeds for reproducibility."""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


# ---------------------------------------------------------------------------
# Training helpers
# ---------------------------------------------------------------------------

def train_one_epoch(
    model: nn.Module,
    loader: DataLoader,
    optimizer: torch.optim.Optimizer,
    criterion: nn.Module,
    device: torch.device,
) -> float:
    """Run one training epoch and return the average loss.

    Follows strict gradient discipline:
    - ``model.train()`` is assumed to be set by the caller.
    - ``optimizer.zero_grad()`` at the start of each iteration.
    - ``loss.backward()`` → ``optimizer.step()``.
    """
    model.train()
    running_loss = 0.0
    n_batches = 0

    for images, labels in loader:
        images = images.to(device)   # (B, 3, 32, 32)
        labels = labels.to(device)   # (B,)

        optimizer.zero_grad()
        logits = model(images)       # (B, 10)
        loss = criterion(logits, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        n_batches += 1

    return running_loss / max(n_batches, 1)


def save_checkpoint(
    path: str,
    model: nn.Module,
    class_names: list[str],
    config: dict,
) -> None:
    """Save model state dict and metadata to *path*."""
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    torch.save(
        {
            "model_state_dict": model.state_dict(),
            "class_names": class_names,
            "config": config,
        },
        path,
    )


# ---------------------------------------------------------------------------
# Overfit smoke test
# ---------------------------------------------------------------------------

def run_overfit_smoke(device: torch.device) -> bool:
    """Overfit a single batch of 16 samples in 200 steps.

    Returns ``True`` if final loss < 0.1, ``False`` otherwise.
    """
    print("[smoke] Running 1-batch overfit test …")

    # Tiny dataset: 16 samples, fixed seed
    dataset = get_synthetic_dataset(n_samples=16, seed=0)
    images, labels = dataset.tensors

    images = images.to(device)  # (16, 3, 32, 32)
    labels = labels.to(device)  # (16,)

    model = TinyCNN(num_classes=10).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
    criterion = nn.CrossEntropyLoss()

    model.train()
    losses = []
    for step in range(200):
        optimizer.zero_grad()
        logits = model(images)
        loss = criterion(logits, labels)
        loss.backward()
        optimizer.step()
        losses.append(loss.item())

    final_loss = losses[-1]
    ok = final_loss < 0.1
    status = "PASS" if ok else "FAIL"
    print(
        f"[smoke] Overfit test {status}: "
        f"initial loss {losses[0]:.4f} → final loss {final_loss:.4f}"
    )
    return ok


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main(args: argparse.Namespace | None = None) -> None:
    parser = argparse.ArgumentParser(description="Train ImageClassifier")
    parser.add_argument("--epochs", type=int, default=10, help="Number of training epochs")
    parser.add_argument("--batch", type=int, default=64, help="Batch size")
    parser.add_argument("--lr", type=float, default=1e-3, help="Learning rate")
    parser.add_argument("--seed", type=int, default=42, help="Random seed")
    parser.add_argument(
        "--test",
        action="store_true",
        help="Run a 1-batch overfit smoke test and exit",
    )
    if args is None:
        args = parser.parse_args()

    seed_everything(args.seed)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"[train] Device: {device}")

    class_names = [f"class_{i}" for i in range(10)]

    # ---- Smoke-test mode ----
    if args.test:
        ok = run_overfit_smoke(device)

        # Also save a checkpoint so inference can be tested
        model = TinyCNN(num_classes=10).to(device)
        # Train it a bit so it's not completely random
        dataset = get_synthetic_dataset(n_samples=64, seed=0)
        loader = DataLoader(dataset, batch_size=64, shuffle=False)
        optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
        criterion = nn.CrossEntropyLoss()
        for _ in range(10):
            model.train()
            for images, labels in loader:
                images, labels = images.to(device), labels.to(device)
                optimizer.zero_grad()
                loss = criterion(model(images), labels)
                loss.backward()
                optimizer.step()

        ckpt_path = "checkpoints/best.pt"
        save_checkpoint(ckpt_path, model, class_names, {
            "epochs": 10,
            "batch": 64,
            "lr": 1e-2,
            "seed": args.seed,
        })
        print(f"[train] Checkpoint saved to {ckpt_path}")

        if not ok:
            print("[train] OVERFIT SMOKE FAILED — model architecture may be broken.")
            sys.exit(1)

        print("[train] All smoke tests passed.")
        sys.exit(0)

    # ---- Normal training mode ----
    n_train = max(args.batch * 4, 256)  # at least 4 batches
    print(f"[train] Generating {n_train} synthetic training samples …")
    train_dataset = get_synthetic_dataset(n_samples=n_train, seed=args.seed)
    train_loader = DataLoader(
        train_dataset,
        batch_size=args.batch,
        shuffle=True,
        pin_memory=device.type == "cuda",
    )

    model = TinyCNN(num_classes=10).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
    criterion = nn.CrossEntropyLoss()

    print(
        f"[train] Starting training: "
        f"{args.epochs} epochs, batch={args.batch}, lr={args.lr}"
    )

    for epoch in range(args.epochs):
        avg_loss = train_one_epoch(model, train_loader, optimizer, criterion, device)
        print(f"  epoch {epoch + 1:>3d}/{args.epochs}  loss = {avg_loss:.4f}")

    # Save best checkpoint
    ckpt_path = "checkpoints/best.pt"
    save_checkpoint(ckpt_path, model, class_names, {
        "epochs": args.epochs,
        "batch": args.batch,
        "lr": args.lr,
        "seed": args.seed,
    })
    print(f"[train] Checkpoint saved to {ckpt_path}")


if __name__ == "__main__":
    main()