| """ |
| train.py -- Part B, Section 3.3 (model training + loss-function ablation). |
| |
| Two arms of the ablation, identical in every other respect: |
| python train.py --loss ce --epochs 15 |
| python train.py --loss focal --epochs 15 --focal-alpha 0.25 --focal-gamma 2.0 |
| python train.py --loss weighted_ce --fg-weight 4.0 --epochs 15 # optional third arm |
| |
| Model selection: best validation mAP@0.50 (not loss), with early stopping. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import math |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| from common import (DEFAULT_IMG_SIZE, make_grad_scaler, autocast, DET_SPLIT_CSV, METRICS_DIR, PLOTS_DIR, |
| SEED, SEED_TAG, WEIGHTS_DIR, KvasirDetectionDataset, _mpl, |
| average_precision, coco_map, get_device, make_loader, |
| progress, save_json, set_seed) |
| from model import build_detector, parameter_report, set_roi_loss |
|
|
|
|
| def run_tag(args) -> str: |
| t = f"det_{args.arch}_{args.loss}" |
| if args.loss == "focal": |
| t += f"-a{args.focal_alpha}g{args.focal_gamma}" |
| if args.tag: |
| t += f"_{args.tag}" |
| return f"{t}_{SEED_TAG}" |
|
|
|
|
| |
| @torch.no_grad() |
| def run_inference(model, loader, device): |
| """Returns (preds, gts) in the flat dict format the metric helpers expect.""" |
| model.eval() |
| preds, gts = [], [] |
| for images, targets in progress(loader, desc="infer"): |
| images = [im.to(device) for im in images] |
| outs = model(images) |
| for o, t in zip(outs, targets): |
| preds.append({"boxes": o["boxes"].detach().cpu().numpy(), |
| "scores": o["scores"].detach().cpu().numpy(), |
| "labels": o["labels"].detach().cpu().numpy()}) |
| gts.append({"boxes": t["boxes"].numpy(), "stem": t["stem"]}) |
| return preds, gts |
|
|
|
|
| def warmup_lr(optimizer, base_lr: float, step: int, warmup_steps: int): |
| if warmup_steps <= 0 or step >= warmup_steps: |
| return |
| f = (step + 1) / warmup_steps |
| for g in optimizer.param_groups: |
| g["lr"] = base_lr * f |
|
|
|
|
| |
| def main(): |
| ap = argparse.ArgumentParser(description="Kvasir-SEG detector training (seed 16)") |
| ap.add_argument("--arch", default="fasterrcnn_resnet50_fpn", |
| choices=["fasterrcnn_resnet50_fpn", "retinanet_resnet50_fpn"]) |
| ap.add_argument("--loss", default="ce", choices=["ce", "focal", "weighted_ce"], |
| help="ROI-head classification loss (the ablation variable)") |
| ap.add_argument("--focal-alpha", type=float, default=0.25) |
| ap.add_argument("--focal-gamma", type=float, default=2.0) |
| ap.add_argument("--fg-weight", type=float, default=4.0) |
| ap.add_argument("--epochs", type=int, default=15) |
| ap.add_argument("--bs", type=int, default=4) |
| ap.add_argument("--lr", type=float, default=5e-3) |
| ap.add_argument("--optimizer", default="sgd", choices=["sgd", "adamw"]) |
| ap.add_argument("--momentum", type=float, default=0.9) |
| ap.add_argument("--wd", type=float, default=5e-4) |
| ap.add_argument("--img-size", type=int, default=DEFAULT_IMG_SIZE) |
| ap.add_argument("--trainable-layers", type=int, default=3, |
| help="ResNet stages left unfrozen (0-5)") |
| ap.add_argument("--no-augment", action="store_true") |
| ap.add_argument("--no-pretrained", action="store_true") |
| ap.add_argument("--warmup-frac", type=float, default=0.1, |
| help="fraction of the first epoch used for linear LR warmup") |
| ap.add_argument("--patience", type=int, default=5) |
| ap.add_argument("--min-delta", type=float, default=1e-4) |
| ap.add_argument("--clip-grad", type=float, default=10.0) |
| ap.add_argument("--workers", type=int, default=2) |
| ap.add_argument("--amp", action="store_true", default=True) |
| ap.add_argument("--tag", default="") |
| args = ap.parse_args() |
|
|
| set_seed(SEED) |
| device = get_device() |
| tag = run_tag(args) |
| loss_desc = set_roi_loss(args.loss, alpha=args.focal_alpha, |
| gamma=args.focal_gamma, fg_weight=args.fg_weight) |
| print(json.dumps({"run": tag, "seed": SEED, "device": str(device), |
| "loss": loss_desc, **vars(args)}, indent=2)) |
|
|
| ds_tr = KvasirDetectionDataset(DET_SPLIT_CSV, "train", augment=not args.no_augment, |
| img_size=args.img_size) |
| ds_va = KvasirDetectionDataset(DET_SPLIT_CSV, "val", augment=False, |
| img_size=args.img_size) |
| print(f"[data] train={len(ds_tr)} val={len(ds_va)} img_size={args.img_size} " |
| f"augment={not args.no_augment}") |
| dl_tr = make_loader(ds_tr, args.bs, shuffle=True, num_workers=args.workers) |
| dl_va = make_loader(ds_va, max(1, args.bs), shuffle=False, num_workers=args.workers) |
|
|
| model = build_detector(args.arch, pretrained_backbone=not args.no_pretrained, |
| img_size=args.img_size, |
| trainable_backbone_layers=args.trainable_layers).to(device) |
| print(f"[model] {args.arch}: {parameter_report(model)}") |
|
|
| params = [p for p in model.parameters() if p.requires_grad] |
| if args.optimizer == "sgd": |
| opt = torch.optim.SGD(params, lr=args.lr, momentum=args.momentum, weight_decay=args.wd) |
| else: |
| opt = torch.optim.AdamW(params, lr=args.lr, weight_decay=args.wd) |
| sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=args.epochs, |
| eta_min=args.lr * 0.01) |
| scaler = make_grad_scaler(enabled=(device.type == "cuda" and args.amp)) |
| warmup_steps = int(args.warmup_frac * max(1, len(dl_tr))) |
|
|
| hist, best, bad, gstep = [], -1.0, 0, 0 |
| ckpt_path = WEIGHTS_DIR / f"{tag}.pt" |
| t0 = time.time() |
|
|
| for ep in range(1, args.epochs + 1): |
| model.train() |
| agg, n = {}, 0 |
| for images, targets in progress(dl_tr, desc=f"ep{ep} train"): |
| images = [im.to(device) for im in images] |
| targets = [{k: (v.to(device) if torch.is_tensor(v) else v) |
| for k, v in t.items()} for t in targets] |
| if gstep < warmup_steps: |
| warmup_lr(opt, args.lr, gstep, warmup_steps) |
| opt.zero_grad(set_to_none=True) |
| with autocast(enabled=scaler.is_enabled()): |
| losses = model(images, targets) |
| total = sum(losses.values()) |
| if not math.isfinite(float(total.detach())): |
| print(f"[warn] non-finite loss at step {gstep}, skipping batch") |
| gstep += 1 |
| continue |
| scaler.scale(total).backward() |
| if args.clip_grad: |
| scaler.unscale_(opt) |
| torch.nn.utils.clip_grad_norm_(params, args.clip_grad) |
| scaler.step(opt); scaler.update() |
| for k, v in losses.items(): |
| agg[k] = agg.get(k, 0.0) + float(v) |
| agg["total"] = agg.get("total", 0.0) + float(total) |
| n += 1 |
| gstep += 1 |
| train_losses = {k: v / max(1, n) for k, v in agg.items()} |
|
|
| preds, gts = run_inference(model, dl_va, device) |
| ap50 = average_precision(preds, gts, 0.5)["ap"] |
| ap75 = average_precision(preds, gts, 0.75)["ap"] |
| sched.step() |
|
|
| row = {"epoch": ep, **{f"loss_{k}": v for k, v in train_losses.items()}, |
| "val_mAP50": float(ap50), "val_mAP75": float(ap75), |
| "lr": opt.param_groups[0]["lr"]} |
| hist.append(row) |
| print(f"[{tag}] ep {ep:03d} | loss {train_losses.get('total', 0):.4f} " |
| f"(cls {train_losses.get('loss_classifier', 0):.4f}, " |
| f"box {train_losses.get('loss_box_reg', 0):.4f}) " |
| f"| val mAP@50 {ap50:.4f} | mAP@75 {ap75:.4f}") |
|
|
| if ap50 > best + args.min_delta: |
| best, bad = ap50, 0 |
| torch.save({"model_state": model.state_dict(), "args": vars(args), |
| "epoch": ep, "val_mAP50": float(ap50), "seed": SEED, |
| "tag": tag, "loss_description": loss_desc}, ckpt_path) |
| print(f" -> new best, saved {ckpt_path.name}") |
| else: |
| bad += 1 |
| if bad >= args.patience: |
| print(f"[{tag}] early stop after {ep} epochs") |
| break |
|
|
| summary = {"tag": tag, "seed": SEED, "best_val_mAP50": float(best), |
| "epochs_run": len(hist), "minutes": round((time.time() - t0) / 60, 2), |
| "loss_description": loss_desc, "config": vars(args), |
| "parameter_report": parameter_report(model), "history": hist, |
| "training_recipe": { |
| "architecture": args.arch, |
| "input_resolution": f"{args.img_size}x{args.img_size} (letterbox-free square resize)", |
| "annotation_format": "xyxy absolute pixels + label 1 (polyp), " |
| "converted from kvasir bbox json", |
| "optimizer": args.optimizer, "base_lr": args.lr, |
| "lr_schedule": f"linear warmup over {warmup_steps} steps, then cosine to " |
| f"{args.lr * 0.01:.2e}", |
| "batch_size": args.bs, "weight_decay": args.wd, |
| "augmentation": "none" if args.no_augment else |
| "random h/v flip + brightness/contrast + saturation jitter", |
| "inference_score_threshold": 0.05, |
| "nms_iou_threshold": 0.5, |
| "stopping_rule": f"early stop on val mAP@0.50, patience={args.patience}, " |
| f"min_delta={args.min_delta}; best checkpoint restored"}} |
| save_json(summary, METRICS_DIR / f"train_{tag}.json") |
|
|
| plt = _mpl() |
| fig, axes = plt.subplots(1, 2, figsize=(11, 4.2)) |
| ep_ = [h["epoch"] for h in hist] |
| for k in [k for k in hist[0] if k.startswith("loss_")]: |
| axes[0].plot(ep_, [h[k] for h in hist], marker="o", ms=3, label=k.replace("loss_", "")) |
| axes[0].set_xlabel("epoch"); axes[0].set_ylabel("loss"); axes[0].legend(fontsize=8) |
| axes[0].set_title(f"{tag}: training losses") |
| axes[1].plot(ep_, [h["val_mAP50"] for h in hist], marker="o", ms=3, label="val mAP@0.50") |
| axes[1].plot(ep_, [h["val_mAP75"] for h in hist], marker="s", ms=3, label="val mAP@0.75") |
| axes[1].set_xlabel("epoch"); axes[1].legend(); axes[1].grid(alpha=0.3) |
| axes[1].set_title(f"{tag}: validation mAP") |
| fig.tight_layout(); fig.savefig(PLOTS_DIR / f"curves_{tag}.png", dpi=160); plt.close(fig) |
| print(f"[saved] {PLOTS_DIR / f'curves_{tag}.png'}") |
| print(f"[done] best val mAP@0.50 = {best:.4f} | weights: {ckpt_path}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
