src/objectmodel_v1/train.py

from __future__ import annotations

import argparse
import json
import os
import time
from pathlib import Path

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel
from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader, DistributedSampler

from .config import apply_overrides, load_config, save_config
from .data import build_dataset, detection_collate
from .evaluate import evaluate_coco
from .losses import ObjectModelCriterion
from .model import build_model
from .utils import (
    ModelEMA,
    learning_rate_factor,
    move_targets,
    save_checkpoint,
    seed_everything,
    trainable_parameter_count,
)


def distributed_context() -> tuple[int, int, int]:
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    rank = int(os.environ.get("RANK", "0"))
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    if world_size > 1:
        if not torch.cuda.is_available():
            raise RuntimeError("Distributed training currently requires CUDA")
        torch.cuda.set_device(local_rank)
        dist.init_process_group(backend="nccl")
    return rank, world_size, local_rank


def build_optimizer(model, config: dict) -> AdamW:
    train = config["train"]
    backbone, other = [], []
    for name, parameter in model.named_parameters():
        if not parameter.requires_grad:
            continue
        (backbone if name.startswith("backbone.") else other).append(parameter)
    return AdamW(
        [
            {"params": other, "lr": float(train["lr"])},
            {"params": backbone, "lr": float(train.get("backbone_lr", train["lr"]))},
        ],
        weight_decay=float(train["weight_decay"]),
    )


def reduce_losses(losses: dict[str, torch.Tensor], world_size: int) -> dict[str, float]:
    values = torch.stack([value.detach() for value in losses.values()])
    if world_size > 1:
        dist.all_reduce(values)
        values /= world_size
    return {name: float(value) for name, value in zip(losses, values, strict=True)}


def main() -> None:
    parser = argparse.ArgumentParser(description="Train ObjectModel-v1")
    parser.add_argument("--config", default="configs/objectmodel_v1.yaml")
    parser.add_argument("--data-root", required=True)
    parser.add_argument("--output", default="outputs/objectmodel_v1")
    parser.add_argument("--resume")
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    parser.add_argument("--set", action="append", default=[])
    args = parser.parse_args()

    rank, world_size, local_rank = distributed_context()
    config = apply_overrides(load_config(args.config), args.set)
    train_config = config["train"]
    seed_everything(int(train_config["seed"]) + rank)
    device = torch.device(f"cuda:{local_rank}" if world_size > 1 else args.device)
    if device.type == "cuda":
        torch.backends.cudnn.benchmark = True
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True
        torch.set_float32_matmul_precision("high")
    channels_last = bool(train_config.get("channels_last", False))
    compile_model = bool(train_config.get("compile", False))
    output_dir = Path(args.output)
    if rank == 0:
        output_dir.mkdir(parents=True, exist_ok=True)
        save_config(config, output_dir / "config.yaml")

    train_dataset = build_dataset(config, args.data_root, "train")
    train_sampler = DistributedSampler(train_dataset, shuffle=True) if world_size > 1 else None
    train_loader = DataLoader(
        train_dataset,
        batch_size=int(train_config["batch_size"]),
        shuffle=train_sampler is None,
        sampler=train_sampler,
        num_workers=int(train_config["workers"]),
        pin_memory=device.type == "cuda",
        drop_last=True,
        persistent_workers=int(train_config["workers"]) > 0,
        prefetch_factor=int(train_config.get("prefetch_factor", 4))
        if int(train_config["workers"]) > 0
        else None,
        collate_fn=detection_collate,
    )
    model = build_model(config).to(device)
    if channels_last:
        model = model.to(memory_format=torch.channels_last)
    criterion = ObjectModelCriterion(config).to(device)
    optimizer = build_optimizer(model, config)
    total_steps = int(train_config["epochs"]) * len(train_loader)
    scheduler = LambdaLR(
        optimizer,
        lambda step: learning_rate_factor(
            step,
            total_steps,
            int(train_config["warmup_steps"]),
            float(train_config["min_lr_ratio"]),
        ),
    )
    ema = ModelEMA(model, float(train_config["ema_decay"])) if rank == 0 else None
    start_epoch, global_step, best_ap = 0, 0, -1.0
    if args.resume:
        checkpoint = torch.load(args.resume, map_location="cpu", weights_only=False)
        model.load_state_dict(checkpoint["model"])
        optimizer.load_state_dict(checkpoint["optimizer"])
        scheduler.load_state_dict(checkpoint["scheduler"])
        start_epoch = int(checkpoint["epoch"]) + 1
        global_step = int(checkpoint.get("global_step", start_epoch * len(train_loader)))
        best_ap = float(checkpoint.get("best_ap", -1.0))
        if ema is not None and "ema" in checkpoint:
            ema.model.load_state_dict(checkpoint["ema"])
    if rank == 0:
        print(
            json.dumps(
                {
                    "parameters": trainable_parameter_count(model),
                    "world_size": world_size,
                    "device": str(device),
                    "steps_per_epoch": len(train_loader),
                },
                indent=2,
            )
        )

    training_model = (
        DistributedDataParallel(model, device_ids=[local_rank], find_unused_parameters=False)
        if world_size > 1
        else model
    )
    if compile_model:
        training_model = torch.compile(training_model, dynamic=False, mode="reduce-overhead")
    use_amp = bool(train_config.get("amp", True)) and device.type == "cuda"
    amp_dtype = (
        torch.bfloat16 if train_config.get("amp_dtype", "float16") == "bfloat16" else torch.float16
    )
    scaler = torch.amp.GradScaler("cuda", enabled=use_amp and amp_dtype == torch.float16)
    history_path = output_dir / "metrics.jsonl"
    for epoch in range(start_epoch, int(train_config["epochs"])):
        if train_sampler is not None:
            train_sampler.set_epoch(epoch)
        training_model.train()
        epoch_start = time.perf_counter()
        log_start = epoch_start
        running = torch.zeros((), device=device)
        for batch_index, (images, targets) in enumerate(train_loader):
            images = images.to(
                device,
                non_blocking=True,
                memory_format=torch.channels_last if channels_last else torch.preserve_format,
            )
            targets = move_targets(targets, device)
            optimizer.zero_grad(set_to_none=True)
            with torch.autocast(device_type=device.type, dtype=amp_dtype, enabled=use_amp):
                outputs = training_model(images)
                losses = criterion(outputs, targets)
            scaler.scale(losses["loss_total"]).backward()
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(
                training_model.parameters(), float(train_config["clip_grad_norm"])
            )
            scaler.step(optimizer)
            scaler.update()
            scheduler.step()
            global_step += 1
            if ema is not None:
                ema.update(model)
            running += losses["loss_total"].detach()
            if rank == 0 and (batch_index + 1) % int(train_config["print_freq"]) == 0:
                now = time.perf_counter()
                log_steps = int(train_config["print_freq"])
                avg_loss = running / (batch_index + 1)
                if world_size > 1:
                    dist.all_reduce(avg_loss)
                    avg_loss = avg_loss / world_size
                print(
                    f"epoch={epoch + 1} step={batch_index + 1}/{len(train_loader)} "
                    f"loss={avg_loss.item():.4f} lr={scheduler.get_last_lr()[0]:.3e} "
                    f"step_seconds={(now - log_start) / log_steps:.3f} "
                    f"images_per_second={log_steps * len(images) / (now - log_start):.2f}",
                    flush=True,
                )
                log_start = now

        epoch_avg_loss = running / max(len(train_loader), 1)
        if world_size > 1:
            dist.all_reduce(epoch_avg_loss)
            epoch_avg_loss = epoch_avg_loss / world_size
        metrics: dict[str, float] = {
            "epoch": epoch + 1,
            "train_loss": epoch_avg_loss.item(),
            "epoch_seconds": time.perf_counter() - epoch_start,
        }
        if world_size > 1:
            dist.barrier()
        should_evaluate = (epoch + 1) % int(train_config["eval_every"]) == 0
        if rank == 0 and should_evaluate:
            val_dataset = build_dataset(config, args.data_root, "val")
            val_loader = DataLoader(
                val_dataset,
                batch_size=int(train_config.get("eval_batch_size", train_config["batch_size"])),
                shuffle=False,
                num_workers=int(train_config["workers"]),
                pin_memory=device.type == "cuda",
                collate_fn=detection_collate,
            )
            metrics.update(
                evaluate_coco(
                    ema.model if ema is not None else model,
                    val_loader,
                    device,
                    output_dir / f"predictions_epoch_{epoch + 1:03d}.json",
                )
            )
        if rank == 0:
            state = {
                "epoch": epoch,
                "global_step": global_step,
                "best_ap": max(best_ap, metrics.get("AP", -1.0)),
                "model": model.state_dict(),
                "ema": ema.model.state_dict() if ema is not None else model.state_dict(),
                "optimizer": optimizer.state_dict(),
                "scheduler": scheduler.state_dict(),
                "config": config,
            }
            save_checkpoint(output_dir / "last.pt", **state)
            if metrics.get("AP", -1.0) > best_ap:
                best_ap = metrics["AP"]
                state["best_ap"] = best_ap
                save_checkpoint(output_dir / "best.pt", **state)
            with history_path.open("a", encoding="utf-8") as handle:
                handle.write(json.dumps(metrics) + "\n")
            print(json.dumps(metrics))
        if world_size > 1:
            dist.barrier()
    if world_size > 1:
        dist.destroy_process_group()


if __name__ == "__main__":
    main()