""" Train MobileNetV3-small on ChestMNIST (14-class multi-label). Outputs: models/chestmnist_mobilenetv3/ mobilenetv3_chestmnist.pth – PyTorch checkpoint (Model A) mobilenetv3_chestmnist.onnx – ONNX export (Model B) training_metrics.json baseline_stats.json – pixel stats for drift detection Usage (quick, shared-server-safe): python3 scripts/train_chestmnist.py --epochs 5 --batch-size 32 """ import argparse import json import logging import time from pathlib import Path import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np import onnx import torch import torch.nn as nn import torch.onnx from PIL import Image from sklearn.metrics import ( average_precision_score, f1_score, roc_auc_score, ) from torch.utils.data import DataLoader from torchvision import transforms from torchvision.models import MobileNet_V3_Small_Weights, mobilenet_v3_small try: import medmnist from medmnist import ChestMNIST, INFO except ImportError as exc: raise SystemExit("medmnist not installed — run: pip install medmnist") from exc BASE_DIR = Path(__file__).resolve().parents[1] OUTPUT_DIR = BASE_DIR / "models" / "chestmnist_mobilenetv3" CHESTMNIST_CLASSES = [ "Atelectasis", "Cardiomegaly", "Effusion", "Infiltration", "Mass", "Nodule", "Pneumonia", "Pneumothorax", "Consolidation", "Edema", "Emphysema", "Fibrosis", "Pleural_Thickening", "Hernia", ] NUM_CLASSES = 14 def get_transforms(image_size: int = 224): train_tf = transforms.Compose([ transforms.Resize((image_size, image_size)), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness=0.2, contrast=0.2), transforms.ToTensor(), transforms.Normalize([0.5] * 3, [0.5] * 3), ]) val_tf = transforms.Compose([ transforms.Resize((image_size, image_size)), transforms.ToTensor(), transforms.Normalize([0.5] * 3, [0.5] * 3), ]) return train_tf, val_tf def build_model(num_classes: int = NUM_CLASSES) -> nn.Module: model = mobilenet_v3_small(weights=MobileNet_V3_Small_Weights.IMAGENET1K_V1) in_features = model.classifier[3].in_features model.classifier[3] = nn.Linear(in_features, num_classes) return model def load_chestmnist(split: str, transform, download: bool, size: int = 224, max_samples: int | None = None): ds = ChestMNIST(split=split, transform=transform, download=download, size=size, as_rgb=True) if max_samples and len(ds) > max_samples: indices = list(range(max_samples)) from torch.utils.data import Subset ds = Subset(ds, indices) return ds def compute_baseline_stats(loader: DataLoader) -> dict: """Compute pixel mean/std from training set for drift detection.""" pixels = [] for images, _ in loader: pixels.append(images.numpy()) if len(pixels) >= 20: # Sample first 20 batches break arr = np.concatenate(pixels, axis=0) # (N, C, H, W) flat = arr.reshape(arr.shape[0], -1) return { "pixel_mean": float(flat.mean()), "pixel_std": float(flat.std()), "channel_means": arr.mean(axis=(0, 2, 3)).tolist(), "channel_stds": arr.std(axis=(0, 2, 3)).tolist(), "n_samples": int(arr.shape[0]), } def tune_thresholds(model: nn.Module, loader: DataLoader, device: torch.device) -> list[float]: """Best-F1 threshold per class on validation set.""" model.eval() all_probs, all_labels = [], [] with torch.no_grad(): for images, labels in loader: images = images.to(device) logits = model(images) probs = torch.sigmoid(logits).cpu().numpy() all_probs.append(probs) all_labels.append(labels.numpy().astype(float)) probs = np.vstack(all_probs) labels = np.vstack(all_labels) thresholds = [] for i in range(NUM_CLASSES): best_t, best_f1 = 0.5, 0.0 for t in np.arange(0.10, 0.90, 0.05): preds = (probs[:, i] >= t).astype(int) f = f1_score(labels[:, i], preds, zero_division=0) if f > best_f1: best_f1, best_t = f, t thresholds.append(round(float(best_t), 3)) return thresholds def evaluate(model: nn.Module, loader: DataLoader, device: torch.device, thresholds: list[float]) -> dict: model.eval() all_probs, all_labels = [], [] with torch.no_grad(): for images, labels in loader: images = images.to(device) logits = model(images) probs = torch.sigmoid(logits).cpu().numpy() all_probs.append(probs) all_labels.append(labels.numpy().astype(float)) probs = np.vstack(all_probs) labels = np.vstack(all_labels) thr_arr = np.array(thresholds) preds = (probs >= thr_arr).astype(int) per_class_auroc, per_class_auprc, per_class_f1 = {}, {}, {} for i, cls in enumerate(CHESTMNIST_CLASSES): if labels[:, i].sum() > 0: per_class_auroc[cls] = round(float(roc_auc_score(labels[:, i], probs[:, i])), 4) per_class_auprc[cls] = round(float(average_precision_score(labels[:, i], probs[:, i])), 4) else: per_class_auroc[cls] = None per_class_auprc[cls] = None per_class_f1[cls] = round(float(f1_score(labels[:, i], preds[:, i], zero_division=0)), 4) macro_auroc_vals = [v for v in per_class_auroc.values() if v is not None] macro_auprc_vals = [v for v in per_class_auprc.values() if v is not None] return { "per_class_auroc": per_class_auroc, "per_class_auprc": per_class_auprc, "per_class_f1": per_class_f1, "test_macro_roc_auc": round(float(np.mean(macro_auroc_vals)), 4) if macro_auroc_vals else None, "test_macro_auprc": round(float(np.mean(macro_auprc_vals)), 4) if macro_auprc_vals else None, "test_micro_f1": round(float(f1_score(labels, preds, average="micro", zero_division=0)), 4), "test_macro_f1": round(float(f1_score(labels, preds, average="macro", zero_division=0)), 4), } def export_onnx(model: nn.Module, output_path: Path, image_size: int, device: torch.device): model.eval() dummy = torch.randn(1, 3, image_size, image_size).to(device) torch.onnx.export( model, dummy, str(output_path), input_names=["input"], output_names=["logits"], dynamic_axes={"input": {0: "batch_size"}, "logits": {0: "batch_size"}}, opset_version=17, ) onnx.checker.check_model(str(output_path)) print(f" ONNX saved → {output_path}") def train_epoch(model: nn.Module, loader: DataLoader, optimizer: torch.optim.Optimizer, criterion: nn.Module, device: torch.device) -> float: model.train() total_loss = 0.0 for images, labels in loader: images, labels = images.to(device), labels.float().to(device) optimizer.zero_grad() loss = criterion(model(images), labels) loss.backward() optimizer.step() total_loss += loss.item() return total_loss / max(len(loader), 1) def save_training_plots(history: list[dict], test_metrics: dict, output_dir: Path): plots_dir = output_dir / "plots" plots_dir.mkdir(parents=True, exist_ok=True) # 1. Learning Curve epochs = [h["epoch"] for h in history] train_loss = [h["train_loss"] for h in history] val_loss = [h["val_loss"] for h in history] plt.figure(figsize=(8, 5)) plt.plot(epochs, train_loss, label="Train Loss", marker="o", color="#2563eb") plt.plot(epochs, val_loss, label="Val Loss", marker="o", color="#dc2626") plt.title("Training & Validation Loss Curve") plt.xlabel("Epoch") plt.ylabel("BCE Loss") plt.grid(True, alpha=0.3) plt.legend() plt.tight_layout() plt.savefig(plots_dir / "loss_curve.png", dpi=150) plt.close() # 2. Per-class metrics bar chart auroc = test_metrics["per_class_auroc"] auprc = test_metrics["per_class_auprc"] labels = [k for k in auroc.keys() if auroc[k] is not None] auroc_vals = [auroc[k] for k in labels] auprc_vals = [auprc[k] for k in labels] y = np.arange(len(labels)) fig, ax = plt.subplots(figsize=(10, 8)) ax.barh(y - 0.2, auroc_vals, height=0.4, label="AUROC", color="#3b82f6") ax.barh(y + 0.2, auprc_vals, height=0.4, label="AUPRC", color="#10b981") ax.set_yticks(y) ax.set_yticklabels(labels, fontweight="bold") ax.set_xlim(0, 1.05) ax.set_title("Per-Class AUROC & AUPRC", fontweight="bold") ax.grid(axis="x", alpha=0.3) ax.legend() plt.tight_layout() plt.savefig(plots_dir / "per_class_metrics.png", dpi=150) plt.close() def main(): parser = argparse.ArgumentParser(description="Train MobileNetV3-small on ChestMNIST (14-class multi-label)") parser.add_argument("--epochs", type=int, default=15) parser.add_argument("--batch-size", type=int, default=32) parser.add_argument("--lr", type=float, default=1e-3) parser.add_argument("--image-size", type=int, default=224) parser.add_argument("--workers", type=int, default=2) parser.add_argument("--max-train-samples", type=int, default=None) parser.add_argument("--max-val-samples", type=int, default=None) parser.add_argument("--max-test-samples", type=int, default=None) parser.add_argument("--output-dir", type=Path, default=OUTPUT_DIR) parser.add_argument("--no-download", action="store_true") args = parser.parse_args() args.output_dir.mkdir(parents=True, exist_ok=True) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Device: {device}") print(f"Output: {args.output_dir}") train_tf, val_tf = get_transforms(args.image_size) download = not args.no_download print("Loading ChestMNIST…") train_ds = load_chestmnist("train", train_tf, download, args.image_size, args.max_train_samples) val_ds = load_chestmnist("val", val_tf, download, args.image_size, args.max_val_samples) test_ds = load_chestmnist("test", val_tf, download, args.image_size, args.max_test_samples) print(f" Train: {len(train_ds)} Val: {len(val_ds)} Test: {len(test_ds)}") train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True) val_loader = DataLoader(val_ds, batch_size=64, shuffle=False, num_workers=args.workers) test_loader = DataLoader(test_ds, batch_size=64, shuffle=False, num_workers=args.workers) # Baseline pixel stats for drift detection print("Computing baseline stats…") baseline_stats = compute_baseline_stats(train_loader) (args.output_dir / "baseline_stats.json").write_text( json.dumps(baseline_stats, indent=2), encoding="utf-8" ) print(f" Mean={baseline_stats['pixel_mean']:.4f} Std={baseline_stats['pixel_std']:.4f}") model = build_model(NUM_CLASSES).to(device) criterion = nn.BCEWithLogitsLoss() optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4) scheduler = torch.optim.lr_scheduler.OneCycleLR( optimizer, max_lr=args.lr, steps_per_epoch=len(train_loader), epochs=args.epochs, ) best_val_loss = float("inf") history = [] print(f"\nTraining {args.epochs} epochs…") for epoch in range(1, args.epochs + 1): t0 = time.time() train_loss = train_epoch(model, train_loader, optimizer, criterion, device) scheduler.step() # Quick val loss model.eval() val_loss = 0.0 with torch.no_grad(): for imgs, lbls in val_loader: imgs, lbls = imgs.to(device), lbls.float().to(device) val_loss += criterion(model(imgs), lbls).item() val_loss /= max(len(val_loader), 1) elapsed = time.time() - t0 print(f" Epoch {epoch}/{args.epochs} — train_loss={train_loss:.4f} val_loss={val_loss:.4f} ({elapsed:.1f}s)") history.append({"epoch": epoch, "train_loss": round(train_loss, 4), "val_loss": round(val_loss, 4)}) if val_loss < best_val_loss: best_val_loss = val_loss torch.save(model.state_dict(), args.output_dir / "mobilenetv3_chestmnist.pth") print(" ✓ checkpoint saved") # Reload best checkpoint model.load_state_dict(torch.load(args.output_dir / "mobilenetv3_chestmnist.pth", map_location=device)) # Threshold tuning on val set print("\nTuning thresholds on validation set…") thresholds = tune_thresholds(model, val_loader, device) print(f" Thresholds: {thresholds}") # Final evaluation on test set print("\nEvaluating on test set…") test_metrics = evaluate(model, test_loader, device, thresholds) # Save training_metrics.json training_metrics = { "architecture": "MobileNetV3-small", "dataset": "ChestMNIST", "class_names": CHESTMNIST_CLASSES, "num_classes": NUM_CLASSES, "epochs": args.epochs, "batch_size": args.batch_size, "learning_rate": args.lr, "image_size": args.image_size, "thresholds": thresholds, "multi_label": True, "best_val_loss": round(best_val_loss, 4), "history": history, **test_metrics, } (args.output_dir / "training_metrics.json").write_text( json.dumps(training_metrics, indent=2), encoding="utf-8" ) # Save Professional Plots print("\nGenerating training & evaluation plots…") save_training_plots(history, test_metrics, args.output_dir) print(f"\n Macro AUROC : {test_metrics['test_macro_roc_auc']}") print(f" Macro AUPRC : {test_metrics['test_macro_auprc']}") print(f" Micro F1 : {test_metrics['test_micro_f1']}") # ONNX Export onnx_path = args.output_dir / "mobilenetv3_chestmnist.onnx" print(f"\nExporting ONNX → {onnx_path}") export_onnx(model, onnx_path, args.image_size, device) # ONNX export report (for backend resolver) onnx_report = { "base_onnx": str(onnx_path.name), "optimized_onnx": str(onnx_path.name), "serving_onnx": str(onnx_path.name), "input_shape": [1, 3, args.image_size, args.image_size], } (args.output_dir / "onnx_export_report.json").write_text( json.dumps(onnx_report, indent=2), encoding="utf-8" ) print(f"\n✅ All artifacts saved to {args.output_dir}") print(" Model A (PyTorch) : mobilenetv3_chestmnist.pth") print(" Model B (ONNX) : mobilenetv3_chestmnist.onnx") if __name__ == "__main__": main()