scripts/train.py

"""Fine-tune RT-DETR on COCO.

Usage:
    python scripts/train.py --config configs/rtdetr_r50_coco.yaml
"""
from __future__ import annotations

import argparse
import math
import os
import random

import numpy as np
import torch
import yaml
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
from transformers import RTDetrForObjectDetection, RTDetrImageProcessor

try:
    from pycocotools.coco import COCO
    from torchvision.datasets import CocoDetection
except ImportError as e:
    raise SystemExit(f"Missing dependency: {e}\nInstall pycocotools and torchvision.") from e


# ---------------------------------------------------------------------------
# Dataset helpers
# ---------------------------------------------------------------------------

def build_coco_dataset(img_dir: str, ann_file: str, processor: RTDetrImageProcessor):
    base = CocoDetection(root=img_dir, annFile=ann_file)

    cat_ids = sorted(base.coco.getCatIds())
    cat_id_to_idx = {cat_id: idx for idx, cat_id in enumerate(cat_ids)}

    class _Wrapped(torch.utils.data.Dataset):
        def __getitem__(self, idx):
            img, targets = base[idx]
            image_id = targets[0]["image_id"] if targets else 0
            annotations = {
                "image_id": image_id,
                "annotations": [
                    {
                        "bbox": t["bbox"],
                        "category_id": cat_id_to_idx[t["category_id"]],
                        "area": t["bbox"][2] * t["bbox"][3],
                        "iscrowd": t.get("iscrowd", 0),
                    }
                    for t in targets
                ],
            }
            encoding = processor(
                images=img,
                annotations=annotations,
                return_tensors="pt",
            )
            result = {}
            for k, v in encoding.items():
                if isinstance(v, torch.Tensor):
                    result[k] = v.squeeze(0)
                elif isinstance(v, list) and len(v) == 1:
                    result[k] = v[0]
                else:
                    result[k] = v
            return result

        def __len__(self):
            return len(base)

    return _Wrapped()


def collate_fn(batch):
    pixel_values = torch.stack([b["pixel_values"] for b in batch])
    labels = [b["labels"] for b in batch]
    return {"pixel_values": pixel_values, "labels": labels}


# ---------------------------------------------------------------------------
# Training
# ---------------------------------------------------------------------------

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"].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

    model_src = resume if resume else cfg["model"]["id"]
    processor = RTDetrImageProcessor.from_pretrained(model_src)
    model = RTDetrForObjectDetection.from_pretrained(
        model_src,
        num_labels=cfg["model"]["num_labels"],
        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"]:
        try:
            model.gradient_checkpointing_enable()
            print("Gradient checkpointing enabled.")
        except ValueError:
            print("Gradient checkpointing not supported by this model, skipping.")

    # Split params: backbone gets a lower LR
    factor = cfg["training"]["optimizer"]["backbone_lr_factor"]
    base_lr = cfg["training"]["optimizer"]["lr"]
    backbone_params, rest_params = [], []
    for name, param in model.named_parameters():
        if "backbone" in name:
            backbone_params.append(param)
        else:
            rest_params.append(param)

    optimizer = AdamW(
        [
            {"params": backbone_params, "lr": base_lr * factor},
            {"params": rest_params,     "lr": base_lr},
        ],
        weight_decay=cfg["training"]["optimizer"]["weight_decay"],
    )

    epochs = cfg["training"]["epochs"]
    warmup_epochs = cfg["training"]["scheduler"]["warmup_epochs"]
    min_lr = cfg["training"]["scheduler"]["min_lr"]

    warmup = LinearLR(optimizer, start_factor=1e-3, end_factor=1.0, total_iters=warmup_epochs)
    cosine = CosineAnnealingLR(optimizer, T_max=epochs - warmup_epochs, eta_min=min_lr)
    scheduler = SequentialLR(optimizer, schedulers=[warmup, cosine], milestones=[warmup_epochs])

    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
    )
    train_loader = DataLoader(
        train_ds,
        batch_size=cfg["training"]["batch_size"],
        shuffle=True,
        num_workers=cfg["data"]["num_workers"],
        collate_fn=collate_fn,
        pin_memory=True,
    )

    save_dir = cfg["training"]["save_dir"]
    os.makedirs(save_dir, exist_ok=True)
    grad_accum = cfg["training"]["grad_accum_steps"]
    clip_norm = cfg["training"]["clip_grad_norm"]

    for epoch in range(start_epoch, epochs + 1):
        model.train()
        running_loss = 0.0
        optimizer.zero_grad()

        for step, batch in enumerate(train_loader, start=1):
            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("cuda", enabled=use_amp, dtype=amp_dtype):
                outputs = model(pixel_values=pixel_values, labels=labels)
                loss = outputs.loss / grad_accum

            scaler.scale(loss).backward()

            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()

            running_loss += loss.item() * grad_accum
            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:
            ckpt_path = os.path.join(save_dir, f"epoch_{epoch:03d}")
            model.save_pretrained(ckpt_path)
            processor.save_pretrained(ckpt_path)
            print(f"Saved checkpoint → {ckpt_path}")


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

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()


if __name__ == "__main__":
    args = parse_args()
    with open(args.config) as f:
        cfg = yaml.safe_load(f)
    train(cfg, resume=args.resume)