from __future__ import annotations import argparse import json import random from dataclasses import asdict, dataclass from pathlib import Path from typing import Dict, List, Tuple import numpy as np import torch import torch.nn as nn import torch.optim as optim from sklearn.metrics import classification_report from torch.utils.data import DataLoader, random_split from torchvision import datasets, models, transforms @dataclass class TrainConfig: data_dir: str = "data/pokemon" output_model_path: str = "models/custom_resnet18.pth" output_metrics_path: str = "reports/custom_model_metrics.json" batch_size: int = 16 num_epochs: int = 8 learning_rate: float = 1e-4 weight_decay: float = 1e-4 val_split: float = 0.2 image_size: int = 224 seed: int = 42 def set_seed(seed: int) -> None: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) def get_device() -> torch.device: return torch.device("cuda" if torch.cuda.is_available() else "cpu") def build_transforms(image_size: int) -> Tuple[transforms.Compose, transforms.Compose]: train_tfms = transforms.Compose( [ transforms.RandomResizedCrop(image_size), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ] ) eval_tfms = transforms.Compose( [ transforms.Resize(int(image_size * 1.14)), transforms.CenterCrop(image_size), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ] ) return train_tfms, eval_tfms def build_dataloaders(config: TrainConfig) -> Tuple[DataLoader, DataLoader, DataLoader, List[str]]: train_tfms, eval_tfms = build_transforms(config.image_size) train_root = Path(config.data_dir) / "train" test_root = Path(config.data_dir) / "test" full_train_dataset = datasets.ImageFolder(train_root, transform=train_tfms) eval_train_dataset = datasets.ImageFolder(train_root, transform=eval_tfms) test_dataset = datasets.ImageFolder(test_root, transform=eval_tfms) num_train = len(full_train_dataset) num_val = int(num_train * config.val_split) num_train_final = num_train - num_val train_subset, val_subset_indices = random_split( full_train_dataset, [num_train_final, num_val], generator=torch.Generator().manual_seed(config.seed), ) # Recreate val subset with eval transforms by reusing indices from split. val_indices = val_subset_indices.indices val_subset = torch.utils.data.Subset(eval_train_dataset, val_indices) train_loader = DataLoader(train_subset, batch_size=config.batch_size, shuffle=True, num_workers=0) val_loader = DataLoader(val_subset, batch_size=config.batch_size, shuffle=False, num_workers=0) test_loader = DataLoader(test_dataset, batch_size=config.batch_size, shuffle=False, num_workers=0) classes = full_train_dataset.classes return train_loader, val_loader, test_loader, classes def create_model(num_classes: int) -> nn.Module: model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT) in_features = model.fc.in_features model.fc = nn.Linear(in_features, num_classes) return model def evaluate(model: nn.Module, loader: DataLoader, device: torch.device) -> Tuple[float, List[int], List[int]]: model.eval() correct = 0 total = 0 preds_all: List[int] = [] labels_all: List[int] = [] with torch.no_grad(): for images, labels in loader: images = images.to(device) labels = labels.to(device) outputs = model(images) _, preds = torch.max(outputs, dim=1) total += labels.size(0) correct += (preds == labels).sum().item() preds_all.extend(preds.cpu().numpy().tolist()) labels_all.extend(labels.cpu().numpy().tolist()) acc = correct / total if total > 0 else 0.0 return acc, preds_all, labels_all def train(config: TrainConfig) -> Dict[str, object]: set_seed(config.seed) device = get_device() train_loader, val_loader, test_loader, classes = build_dataloaders(config) model = create_model(num_classes=len(classes)).to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=config.learning_rate, weight_decay=config.weight_decay) scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=4, gamma=0.5) best_val_acc = 0.0 best_state = None history = [] for epoch in range(config.num_epochs): model.train() running_loss = 0.0 running_correct = 0 running_total = 0 for images, labels in train_loader: images = images.to(device) labels = labels.to(device) optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() * labels.size(0) _, preds = torch.max(outputs, dim=1) running_correct += (preds == labels).sum().item() running_total += labels.size(0) scheduler.step() train_loss = running_loss / max(1, running_total) train_acc = running_correct / max(1, running_total) val_acc, _, _ = evaluate(model, val_loader, device) history.append( { "epoch": epoch + 1, "train_loss": round(train_loss, 5), "train_acc": round(train_acc, 5), "val_acc": round(val_acc, 5), } ) print( f"Epoch {epoch + 1}/{config.num_epochs} " f"- train_loss: {train_loss:.4f} " f"- train_acc: {train_acc:.4f} " f"- val_acc: {val_acc:.4f}" ) if val_acc > best_val_acc: best_val_acc = val_acc best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()} if best_state is not None: model.load_state_dict(best_state) test_acc, test_preds, test_labels = evaluate(model, test_loader, device) target_names = classes cls_report = classification_report( test_labels, test_preds, target_names=target_names, output_dict=True, zero_division=0, ) output_model = Path(config.output_model_path) output_model.parent.mkdir(parents=True, exist_ok=True) torch.save( { "state_dict": model.state_dict(), "labels": classes, "image_size": config.image_size, "architecture": "resnet18", }, output_model, ) result = { "config": asdict(config), "device": str(device), "num_classes": len(classes), "labels": classes, "best_val_acc": round(best_val_acc, 5), "test_acc": round(test_acc, 5), "history": history, "classification_report": cls_report, "model_path": str(output_model), } output_metrics = Path(config.output_metrics_path) output_metrics.parent.mkdir(parents=True, exist_ok=True) output_metrics.write_text(json.dumps(result, indent=2), encoding="utf-8") print(f"Saved model to: {output_model}") print(f"Saved metrics to: {output_metrics}") print(f"Final test accuracy: {test_acc:.4f}") return result def parse_args() -> TrainConfig: parser = argparse.ArgumentParser(description="Train transfer learning classifier on custom image data.") parser.add_argument("--data-dir", default="data/pokemon") parser.add_argument("--output-model", default="models/custom_resnet18.pth") parser.add_argument("--output-metrics", default="reports/custom_model_metrics.json") parser.add_argument("--batch-size", type=int, default=16) parser.add_argument("--epochs", type=int, default=8) parser.add_argument("--lr", type=float, default=1e-4) parser.add_argument("--weight-decay", type=float, default=1e-4) parser.add_argument("--val-split", type=float, default=0.2) parser.add_argument("--image-size", type=int, default=224) parser.add_argument("--seed", type=int, default=42) args = parser.parse_args() return TrainConfig( data_dir=args.data_dir, output_model_path=args.output_model, output_metrics_path=args.output_metrics, batch_size=args.batch_size, num_epochs=args.epochs, learning_rate=args.lr, weight_decay=args.weight_decay, val_split=args.val_split, image_size=args.image_size, seed=args.seed, ) if __name__ == "__main__": cfg = parse_args() train(cfg)