fix: 10 critical bugs + 5 warnings from pre-training audit

- train.py: add --resume flag, fix grad accum flush, fix COCO annotations format
- evaluate.py: fix label2cat unhashable dict crash
- api/main.py: replace deprecated on_event with lifespan pattern
- train.py + evaluate.py: migrate to torch.amp (PyTorch 2.x)
- configs: A100 tuning — bf16, batch 16, grad_accum 4, effective batch 64
- notebook: fix threshold kwarg, box dict access, nested config keys, checkpoint glob, eval CLI args

api/main.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import io
 import os
+from contextlib import asynccontextmanager
 
 from fastapi import FastAPI, File, HTTPException, Query, UploadFile
 from PIL import Image, UnidentifiedImageError
@@ -12,8 +13,6 @@ from api.schemas import Detection, BoundingBox, PredictResponse
 MODEL_ID = os.getenv("MODEL_ID", "PekingU/rtdetr_r50vd")
 CONFIDENCE_THRESHOLD = float(os.getenv("CONFIDENCE_THRESHOLD", "0.5"))
 
-app = FastAPI(title="detrflow", description="RT-DETR object detection API", version="0.1.0")
-
 _predictor: RTDetrPredictor | None = None
 
 
@@ -27,9 +26,13 @@ def get_predictor() -> RTDetrPredictor:
     return _predictor
 
 
-@app.on_event("startup")
-async def _warm_up() -> None:
+@asynccontextmanager
+async def lifespan(app: FastAPI):
     get_predictor()
+    yield
+
+
+app = FastAPI(title="detrflow", description="RT-DETR object detection API", version="0.1.0", lifespan=lifespan)
 
 
 @app.post("/predict", response_model=PredictResponse, summary="Detect objects in an image")

configs/rtdetr_r50_coco.yaml

@@ -7,12 +7,12 @@ data:
   val_ann:   data/coco/annotations/instances_val2017.json
   train_img: data/coco/train2017
   val_img:   data/coco/val2017
-  num_workers: 4
+  num_workers: 8
 
 training:
   epochs: 12
-  batch_size: 8           # fits 8 GB VRAM with fp16 + gradient checkpointing
-  grad_accum_steps: 2     # effective batch = 16
+  batch_size: 16          # A100 40 GB with bf16 + gradient checkpointing
+  grad_accum_steps: 4     # effective batch = 64 (matches RT-DETR paper)
 
   optimizer:
     name: AdamW
@@ -25,7 +25,8 @@ training:
     warmup_epochs: 1
     min_lr: 1.0e-6
 
-  fp16: true
+  fp16: false
+  bf16: true              # A100 native bfloat16 — wider range, no GradScaler needed
   gradient_checkpointing: true
   clip_grad_norm: 0.1
 

notebooks/detrflow_training.ipynb

@@ -171,27 +171,7 @@
    "id": "show-config",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "import yaml\n",
-    "\n",
-    "with open('configs/rtdetr_r50_coco.yaml') as f:\n",
-    "    cfg = yaml.safe_load(f)\n",
-    "\n",
-    "# Override output dir to point to Drive\n",
-    "cfg['output_dir'] = CKPT_DIR\n",
-    "\n",
-    "# Colab T4 optimizations\n",
-    "cfg['batch_size']       = 8   # T4 16 GB → batch 8 safe with fp16\n",
-    "cfg['grad_accum_steps'] = 2   # effective batch = 16\n",
-    "cfg['num_epochs']       = 12  # standard RT-DETR schedule\n",
-    "cfg['fp16']             = True\n",
-    "\n",
-    "# Save overridden config\n",
-    "with open('configs/rtdetr_r50_coco_colab.yaml', 'w') as f:\n",
-    "    yaml.dump(cfg, f)\n",
-    "\n",
-    "print(yaml.dump(cfg, default_flow_style=False))"
-   ]
+   "source": "import yaml\n\nwith open('configs/rtdetr_r50_coco.yaml') as f:\n    cfg = yaml.safe_load(f)\n\n# Override output dir to point to Drive\ncfg['training']['save_dir'] = CKPT_DIR\n\n# Colab T4 optimizations (nested keys matching train.py's cfg structure)\ncfg['training']['batch_size']       = 8   # T4 16 GB → batch 8 safe with fp16\ncfg['training']['grad_accum_steps'] = 2   # effective batch = 16\ncfg['training']['epochs']           = 12  # standard RT-DETR schedule\ncfg['training']['fp16']             = True\n\n# Save overridden config\nwith open('configs/rtdetr_r50_coco_colab.yaml', 'w') as f:\n    yaml.dump(cfg, f)\n\nprint(yaml.dump(cfg, default_flow_style=False))"
   },
   {
    "cell_type": "markdown",
@@ -209,32 +189,7 @@
    "id": "smoke-test",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "import sys\n",
-    "sys.path.insert(0, '/content/detrflow')\n",
-    "\n",
-    "from inference.predictor import RTDetrPredictor\n",
-    "from inference.visualizer import draw_detections\n",
-    "from PIL import Image\n",
-    "import requests\n",
-    "from IPython.display import display\n",
-    "\n",
-    "# Load sample image\n",
-    "url = 'http://images.cocodataset.org/val2017/000000039769.jpg'\n",
-    "img = Image.open(requests.get(url, stream=True).raw)\n",
-    "\n",
-    "# Run inference\n",
-    "predictor = RTDetrPredictor(model_id='PekingU/rtdetr_r50vd', threshold=0.5)\n",
-    "detections = predictor.predict(img)\n",
-    "\n",
-    "print(f'Detections: {len(detections)}')\n",
-    "for d in detections[:5]:\n",
-    "    print(f\"  {d['label']:20s} score={d['score']:.3f}  box={[round(x) for x in d['box']]}\")\n",
-    "\n",
-    "# Visualize\n",
-    "vis = draw_detections(img, detections)\n",
-    "display(vis)"
-   ]
+   "source": "import sys\nsys.path.insert(0, '/content/detrflow')\n\nfrom inference.predictor import RTDetrPredictor\nfrom inference.visualizer import draw_detections\nfrom PIL import Image\nimport requests\nfrom IPython.display import display\n\n# Load sample image\nurl = 'http://images.cocodataset.org/val2017/000000039769.jpg'\nimg = Image.open(requests.get(url, stream=True).raw)\n\n# Run inference\npredictor = RTDetrPredictor(model_id='PekingU/rtdetr_r50vd', confidence_threshold=0.5)\ndetections = predictor.predict(img)\n\nprint(f'Detections: {len(detections)}')\nfor d in detections[:5]:\n    b = d['box']\n    print(f\"  {d['label']:20s} score={d['score']:.3f}  box=[{b['x1']:.0f},{b['y1']:.0f},{b['x2']:.0f},{b['y2']:.0f}]\")\n\n# Visualize\nvis = draw_detections(img, detections)\ndisplay(vis)"
   },
   {
    "cell_type": "markdown",
@@ -274,24 +229,7 @@
    "id": "resume-training",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "import os\n",
-    "\n",
-    "# Find the latest checkpoint\n",
-    "ckpts = sorted([\n",
-    "    f for f in os.listdir(CKPT_DIR)\n",
-    "    if f.startswith('checkpoint-epoch')\n",
-    "])\n",
-    "\n",
-    "if ckpts:\n",
-    "    latest = f'{CKPT_DIR}/{ckpts[-1]}'\n",
-    "    print(f'Resuming from: {latest}')\n",
-    "    !python scripts/train.py \\\n",
-    "        --config configs/rtdetr_r50_coco_colab.yaml \\\n",
-    "        --resume {latest}\n",
-    "else:\n",
-    "    print('No checkpoint found. Run cell 6 first.')"
-   ]
+   "source": "import os\n\n# Find the latest checkpoint (train.py saves as epoch_001, epoch_002, ...)\nckpts = sorted([\n    f for f in os.listdir(CKPT_DIR)\n    if f.startswith('epoch_')\n])\n\nif ckpts:\n    latest = f'{CKPT_DIR}/{ckpts[-1]}'\n    print(f'Resuming from: {latest}')\n    !python scripts/train.py \\\n        --config configs/rtdetr_r50_coco_colab.yaml \\\n        --resume {latest}\nelse:\n    print('No checkpoint found. Run cell 6 first.')"
   },
   {
    "cell_type": "markdown",
@@ -307,25 +245,7 @@
    "id": "run-eval",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "import os\n",
-    "\n",
-    "# Pick the latest checkpoint (last epoch)\n",
-    "ckpts = sorted([\n",
-    "    f for f in os.listdir(CKPT_DIR)\n",
-    "    if f.startswith('checkpoint-epoch')\n",
-    "])\n",
-    "\n",
-    "if not ckpts:\n",
-    "    print('No checkpoint found. Run training first.')\n",
-    "else:\n",
-    "    best = f'{CKPT_DIR}/{ckpts[-1]}'\n",
-    "    print(f'Evaluating: {best}')\n",
-    "    !python scripts/evaluate.py \\\n",
-    "        --model_path {best} \\\n",
-    "        --coco_dir data/coco \\\n",
-    "        --split val"
-   ]
+   "source": "import os\n\n# Pick the latest checkpoint (train.py saves as epoch_001, epoch_002, ...)\nckpts = sorted([\n    f for f in os.listdir(CKPT_DIR)\n    if f.startswith('epoch_')\n])\n\nif not ckpts:\n    print('No checkpoint found. Run training first.')\nelse:\n    best = f'{CKPT_DIR}/{ckpts[-1]}'\n    print(f'Evaluating: {best}')\n    !python scripts/evaluate.py \\\n        --config configs/rtdetr_r50_coco_colab.yaml \\\n        --checkpoint {best}"
   },
   {
    "cell_type": "markdown",
@@ -385,4 +305,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}

scripts/evaluate.py

@@ -13,7 +13,7 @@ import tempfile
 
 import torch
 import yaml
-from torch.cuda.amp import autocast
+from torch.amp import autocast
 from torch.utils.data import DataLoader
 from transformers import RTDetrForObjectDetection, RTDetrImageProcessor
 
@@ -62,12 +62,15 @@ def build_val_loader(img_dir: str, ann_file: str, processor, batch_size: int, nu
 @torch.inference_mode()
 def evaluate(cfg: dict, checkpoint: str | None = None) -> None:
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    use_fp16 = cfg["training"]["fp16"] and device.type == "cuda"
+    use_fp16 = cfg["training"].get("fp16", False) and device.type == "cuda"
+    use_bf16 = cfg["training"].get("bf16", False) and device.type == "cuda"
+    use_amp = use_fp16 or use_bf16
+    amp_dtype = torch.bfloat16 if use_bf16 else torch.float16
 
     model_id = checkpoint or cfg["model"]["id"]
     processor = RTDetrImageProcessor.from_pretrained(model_id)
     model = RTDetrForObjectDetection.from_pretrained(
-        model_id, torch_dtype=torch.float16 if use_fp16 else torch.float32
+        model_id, torch_dtype=amp_dtype if use_amp else torch.float32
     ).to(device)
     model.eval()
 
@@ -81,14 +84,14 @@ def evaluate(cfg: dict, checkpoint: str | None = None) -> None:
 
     coco_gt = COCO(cfg["data"]["val_ann"])
     # Build HF label → COCO category_id mapping
-    label2cat = {v: k for k, v in coco_gt.cats.items()}
+    label2cat = {cat["name"]: cat["id"] for cat in coco_gt.cats.values()}
 
     results = []
     for batch in loader:
         pixel_values = batch["pixel_values"].to(device)
         orig_sizes = batch["orig_sizes"].to(device)
 
-        with autocast(enabled=use_fp16):
+        with autocast("cuda", enabled=use_amp, dtype=amp_dtype):
             outputs = model(pixel_values=pixel_values)
 
         preds = processor.post_process_object_detection(

scripts/train.py

@@ -13,7 +13,7 @@ import random
 import numpy as np
 import torch
 import yaml
-from torch.cuda.amp import GradScaler, autocast
+from torch.amp import GradScaler, autocast
 from torch.optim import AdamW
 from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR
 from torch.utils.data import DataLoader
@@ -36,14 +36,14 @@ def build_coco_dataset(img_dir: str, ann_file: str, processor: RTDetrImageProces
     class _Wrapped(torch.utils.data.Dataset):
         def __getitem__(self, idx):
             img, targets = base[idx]
-            annotations = [
-                {
-                    "bbox": t["bbox"],           # [x, y, w, h]
-                    "category_id": t["category_id"],
-                    "image_id": t["image_id"],
-                }
-                for t in targets
-            ]
+            image_id = targets[0]["image_id"] if targets else 0
+            annotations = {
+                "image_id": image_id,
+                "annotations": [
+                    {"bbox": t["bbox"], "category_id": t["category_id"]}
+                    for t in targets
+                ],
+            }
             encoding = processor(
                 images=img,
                 annotations=annotations,
@@ -67,22 +67,35 @@ def collate_fn(batch):
 # Training
 # ---------------------------------------------------------------------------
 
-def train(cfg: dict) -> None:
+def train(cfg: dict, resume: str | None = None) -> None:
     seed = cfg["training"]["seed"]
     random.seed(seed)
     np.random.seed(seed)
     torch.manual_seed(seed)
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    use_fp16: bool = cfg["training"]["fp16"] and device.type == "cuda"
+    use_fp16: bool = cfg["training"].get("fp16", False) and device.type == "cuda"
+    use_bf16: bool = cfg["training"].get("bf16", False) and device.type == "cuda"
+    use_amp: bool = use_fp16 or use_bf16
+    amp_dtype = torch.bfloat16 if use_bf16 else torch.float16
 
-    processor = RTDetrImageProcessor.from_pretrained(cfg["model"]["id"])
+    model_src = resume if resume else cfg["model"]["id"]
+    processor = RTDetrImageProcessor.from_pretrained(model_src)
     model = RTDetrForObjectDetection.from_pretrained(
-        cfg["model"]["id"],
+        model_src,
         num_labels=cfg["model"]["num_labels"],
-        ignore_mismatched_sizes=True,
+        ignore_mismatched_sizes=(resume is None),
     ).to(device)
 
+    # Infer which epoch to start from when resuming
+    start_epoch = 1
+    if resume:
+        import re as _re
+        m = _re.search(r"epoch_(\d+)", os.path.basename(resume.rstrip("/\\")))
+        if m:
+            start_epoch = int(m.group(1)) + 1
+        print(f"Resuming from {resume}, starting at epoch {start_epoch}")
+
     if cfg["training"]["gradient_checkpointing"]:
         model.gradient_checkpointing_enable()
 
@@ -112,7 +125,7 @@ def train(cfg: dict) -> None:
     cosine = CosineAnnealingLR(optimizer, T_max=epochs - warmup_epochs, eta_min=min_lr)
     scheduler = SequentialLR(optimizer, schedulers=[warmup, cosine], milestones=[warmup_epochs])
 
-    scaler = GradScaler(enabled=use_fp16)
+    scaler = GradScaler("cuda", enabled=use_fp16)  # GradScaler only needed for fp16, not bf16
 
     train_ds = build_coco_dataset(
         cfg["data"]["train_img"], cfg["data"]["train_ann"], processor
@@ -131,7 +144,7 @@ def train(cfg: dict) -> None:
     grad_accum = cfg["training"]["grad_accum_steps"]
     clip_norm = cfg["training"]["clip_grad_norm"]
 
-    for epoch in range(1, epochs + 1):
+    for epoch in range(start_epoch, epochs + 1):
         model.train()
         running_loss = 0.0
         optimizer.zero_grad()
@@ -140,7 +153,7 @@ def train(cfg: dict) -> None:
             pixel_values = batch["pixel_values"].to(device)
             labels = [{k: v.to(device) for k, v in lbl.items()} for lbl in batch["labels"]]
 
-            with autocast(enabled=use_fp16):
+            with autocast("cuda", enabled=use_amp, dtype=amp_dtype):
                 outputs = model(pixel_values=pixel_values, labels=labels)
                 loss = outputs.loss / grad_accum
 
@@ -157,6 +170,14 @@ def train(cfg: dict) -> None:
             if step % 50 == 0:
                 print(f"[epoch {epoch}/{epochs} step {step}] loss={running_loss/step:.4f}")
 
+        # Flush any remaining accumulated gradients at end of epoch
+        if step % grad_accum != 0:
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), clip_norm)
+            scaler.step(optimizer)
+            scaler.update()
+            optimizer.zero_grad()
+
         scheduler.step()
 
         if epoch % cfg["training"]["save_every_n_epochs"] == 0:
@@ -173,6 +194,7 @@ def train(cfg: dict) -> None:
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser()
     parser.add_argument("--config", default="configs/rtdetr_r50_coco.yaml")
+    parser.add_argument("--resume", default=None, help="Path to checkpoint dir to resume training from")
     return parser.parse_args()
 
 
@@ -180,4 +202,4 @@ if __name__ == "__main__":
     args = parse_args()
     with open(args.config) as f:
         cfg = yaml.safe_load(f)
-    train(cfg)
+    train(cfg, resume=args.resume)