train.py

"""Train a teacher CNN, then compare PTQ INT8, magnitude pruning, and KD distillation."""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import torch

from src.compress import magnitude_prune_filters, model_size_mb, quantize_dynamic_int8
from src.data import get_loaders
from src.engine import benchmark_latency, distillation_loss, evaluate, train_one_epoch
from src.model import StudentCNN, TeacherCNN, count_params

SEED = 42


def set_seed(seed: int) -> None:
    np.random.seed(seed)
    torch.manual_seed(seed)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", choices=["mnist", "cifar10"], default="mnist")
    parser.add_argument("--epochs-teacher", type=int, default=10)
    parser.add_argument("--epochs-finetune", type=int, default=3)
    parser.add_argument("--epochs-distill", type=int, default=10)
    parser.add_argument("--prune-amount", type=float, default=0.5)
    parser.add_argument("--alpha-kd", type=float, default=0.5)
    parser.add_argument("--T-kd", type=float, default=4.0)
    parser.add_argument("--batch-size", type=int, default=128)
    parser.add_argument("--lr", type=float, default=1e-3)
    parser.add_argument("--out", type=str, default="runs")
    args = parser.parse_args()

    set_seed(SEED)
    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    train_loader, test_loader, in_ch = get_loaders(args.dataset, batch_size=args.batch_size)
    sample, _ = next(iter(test_loader))

    print(f"Dataset: {args.dataset}, in_ch={in_ch}")
    print("=== Training teacher ===")
    teacher = TeacherCNN(in_ch=in_ch)
    teacher_params = count_params(teacher)
    print(f"Teacher params: {teacher_params/1e3:.1f} k")
    opt = torch.optim.Adam(teacher.parameters(), lr=args.lr)
    teacher_history = []
    for epoch in range(args.epochs_teacher):
        loss = train_one_epoch(teacher, train_loader, opt)
        acc = evaluate(teacher, test_loader)
        teacher_history.append({"epoch": epoch + 1, "loss": loss, "test_acc": acc})
        print(f"  epoch {epoch+1}: loss={loss:.4f} test_acc={acc*100:.2f}%")
    torch.save(teacher.state_dict(), out_dir / "teacher.pt")

    teacher_acc = evaluate(teacher, test_loader)
    teacher_size = model_size_mb(teacher)
    teacher_lat = benchmark_latency(teacher, sample)
    print(f"Teacher: acc={teacher_acc*100:.2f}% size={teacher_size:.2f} MB lat={teacher_lat:.2f} ms")

    print("=== PTQ dynamic INT8 ===")
    qmodel = quantize_dynamic_int8(teacher)
    q_acc = evaluate(qmodel, test_loader)
    q_size = model_size_mb(qmodel)
    q_lat = benchmark_latency(qmodel, sample)
    print(f"INT8: acc={q_acc*100:.2f}% size={q_size:.2f} MB lat={q_lat:.2f} ms")

    print(f"=== Magnitude pruning ({int(args.prune_amount*100)}%) + fine-tune ===")
    pmodel = magnitude_prune_filters(teacher, amount=args.prune_amount)
    print(f"Pruned (pre-FT): acc={evaluate(pmodel, test_loader)*100:.2f}%")
    opt_p = torch.optim.Adam(pmodel.parameters(), lr=args.lr / 5)
    for epoch in range(args.epochs_finetune):
        loss = train_one_epoch(pmodel, train_loader, opt_p)
        acc = evaluate(pmodel, test_loader)
        print(f"  FT epoch {epoch+1}: loss={loss:.4f} test_acc={acc*100:.2f}%")
    p_acc = evaluate(pmodel, test_loader)
    p_size = model_size_mb(pmodel)
    p_lat = benchmark_latency(pmodel, sample)
    print(f"Pruned: acc={p_acc*100:.2f}% size={p_size:.2f} MB lat={p_lat:.2f} ms")

    print("=== Knowledge distillation to a smaller student ===")
    student = StudentCNN(in_ch=in_ch)
    student_params = count_params(student)
    print(f"Student params: {student_params/1e3:.1f} k")
    teacher.eval()
    opt_s = torch.optim.Adam(student.parameters(), lr=args.lr)
    student_baseline = StudentCNN(in_ch=in_ch)  # same arch, trained without KD for comparison
    opt_b = torch.optim.Adam(student_baseline.parameters(), lr=args.lr)
    distill_history = []
    for epoch in range(args.epochs_distill):
        student.train()
        for x, y in train_loader:
            with torch.no_grad():
                t_logits = teacher(x)
            s_logits = student(x)
            loss = distillation_loss(s_logits, t_logits, y, T=args.T_kd, alpha=args.alpha_kd)
            opt_s.zero_grad()
            loss.backward()
            opt_s.step()
        baseline_loss = train_one_epoch(student_baseline, train_loader, opt_b)
        s_acc = evaluate(student, test_loader)
        b_acc = evaluate(student_baseline, test_loader)
        distill_history.append({"epoch": epoch + 1, "student_kd_acc": s_acc,
                                "student_baseline_acc": b_acc, "baseline_loss": baseline_loss})
        print(f"  KD epoch {epoch+1}: student_kd={s_acc*100:.2f}% baseline={b_acc*100:.2f}%")
    s_acc = evaluate(student, test_loader)
    b_acc = evaluate(student_baseline, test_loader)
    s_size = model_size_mb(student)
    s_lat = benchmark_latency(student, sample)
    print(f"Student (KD): acc={s_acc*100:.2f}% size={s_size:.2f} MB lat={s_lat:.2f} ms")
    print(f"Student (baseline, no KD): acc={b_acc*100:.2f}%")

    variants = [
        {"name": "Teacher (fp32)", "params": teacher_params, "acc": teacher_acc,
         "size_mb": teacher_size, "latency_ms": teacher_lat},
        {"name": "PTQ INT8 (dyn)", "params": teacher_params, "acc": q_acc,
         "size_mb": q_size, "latency_ms": q_lat},
        {"name": f"Pruned {int(args.prune_amount*100)}% + FT", "params": teacher_params,
         "acc": p_acc, "size_mb": p_size, "latency_ms": p_lat},
        {"name": "Student (KD, T=4)", "params": student_params, "acc": s_acc,
         "size_mb": s_size, "latency_ms": s_lat},
        {"name": "Student (no KD)", "params": student_params, "acc": b_acc,
         "size_mb": s_size, "latency_ms": s_lat},
    ]
    metrics = {
        "dataset": args.dataset,
        "teacher_history": teacher_history,
        "distill_history": distill_history,
        "variants": variants,
        "config": vars(args),
    }
    with open(out_dir / "metrics.json", "w", encoding="utf-8") as f:
        json.dump(metrics, f, indent=2)
    print(f"Saved metrics to {out_dir / 'metrics.json'}")


if __name__ == "__main__":
    main()