scripts/train.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import random
from pathlib import Path
import sys

import numpy as np
import torch
import yaml
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler

# ``model`` and ``scripts`` are namespace packages; no __init__.py is required.
PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT))
from scripts.data_loader import MRMSDataset
from model.model_factory import build_model, load_checkpoint


def clone_state_dict(model: torch.nn.Module) -> dict[str, torch.Tensor]:
    """Copy the current parameters to CPU so later epochs cannot mutate them."""
    network = model.module if hasattr(model, "module") else model
    return {name: tensor.detach().cpu().clone() for name, tensor in network.state_dict().items()}


def main() -> None:
    parser = argparse.ArgumentParser(description="NowcastNet single-card training")
    parser.add_argument("--config", default=str(Path(__file__).parents[1] / "conf/config.yaml"))
    parser.add_argument("--data-dir", help="MRMS event directory; defaults to data.data_dir in config.yaml")
    parser.add_argument("--device", default="auto")
    parser.add_argument("--epochs", type=int)
    parser.add_argument("--checkpoint")
    parser.add_argument("--output-dir")
    parser.add_argument("--log-dir", help="training log directory; defaults to training.output_dir in config.yaml")
    parser.add_argument("--checkpoint-prefix")
    parser.add_argument("--height", type=int)
    parser.add_argument("--width", type=int)
    parser.add_argument("--ngf", type=int)
    parser.add_argument("--distributed", action="store_true", help="use torchrun for single-node or multi-node DDP")
    args = parser.parse_args()
    cfg = yaml.safe_load(Path(args.config).read_text())
    mc, dc, tc = cfg["model"], cfg["data"], cfg["training"]
    if args.height: mc["img_height"] = dc["image_height"] = args.height
    if args.width: mc["img_width"] = dc["image_width"] = args.width
    if args.ngf: mc["ngf"] = args.ngf
    data_dir = Path(args.data_dir) if args.data_dir else PROJECT_ROOT / dc["data_dir"]
    distributed = args.distributed or int(__import__("os").environ.get("WORLD_SIZE", "1")) > 1
    rank = 0
    if distributed:
        import os, torch.distributed as dist
        dist.init_process_group(backend="nccl" if torch.cuda.is_available() else "gloo", init_method="env://")
        rank = dist.get_rank()
        if args.device == "auto" and torch.cuda.is_available():
            torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0)))
    if distributed and torch.cuda.is_available():
        device = torch.device("cuda", int(__import__("os").environ.get("LOCAL_RANK", 0)))
    else:
        device = torch.device("cuda" if args.device == "auto" and torch.cuda.is_available() else ("cpu" if args.device == "auto" else args.device))
    seed = int(tc["seed"]); random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    model = build_model(mc, device)
    best_loss = float("inf")
    best_epoch = 0
    best_state = None
    if args.checkpoint:
        checkpoint = load_checkpoint(model, args.checkpoint, device)
        if "best_loss" in checkpoint:
            best_loss = float(checkpoint["best_loss"])
            best_epoch = int(checkpoint.get("epoch", 0))
            best_state = clone_state_dict(model)
    dataset = MRMSDataset(data_dir, dc["image_height"], dc["image_width"], dc["total_length"], "train")
    if not dataset:
        raise RuntimeError(f"No complete MRMS events found under {data_dir}")
    sampler = DistributedSampler(dataset, shuffle=True) if distributed else None
    loader = DataLoader(dataset, batch_size=int(dc["batch_size"]), shuffle=sampler is None, sampler=sampler, num_workers=int(dc["num_workers"]), drop_last=False)
    if distributed:
        from torch.nn.parallel import DistributedDataParallel as DDP
        # The official generator keeps architectural branches whose
        # parameters are not touched by the L1 smoke-training path.
        model = DDP(
            model,
            device_ids=[device.index] if device.type == "cuda" else None,
            find_unused_parameters=True,
        )
    optimizer = torch.optim.Adam(model.parameters(), lr=float(tc["lr"]), weight_decay=float(tc["weight_decay"]))
    epochs = args.epochs or int(tc["epochs"])
    checkpoint_dir = Path(args.output_dir) if args.output_dir else PROJECT_ROOT / mc["checkpoint_dir"]
    checkpoint_prefix = args.checkpoint_prefix or mc.get("checkpoint_prefix", "model_bak")
    log_dir = Path(args.log_dir) if args.log_dir else PROJECT_ROOT / tc["output_dir"]
    loss_history_path = log_dir / tc.get("loss_history_file", "loss_history.json")
    loss_history: list[dict[str, float | int]] = []
    checkpoint_dir.mkdir(parents=True, exist_ok=True)
    if rank == 0:
        log_dir.mkdir(parents=True, exist_ok=True)
    for epoch in range(epochs):
        if sampler is not None: sampler.set_epoch(epoch)
        model.train()
        loss_sum = torch.zeros(1, device=device)
        sample_count = torch.zeros(1, device=device)
        for batch in loader:
            frames = batch["radar_frames"].to(device=device, dtype=torch.float32)
            pred = model(frames)
            target = frames[:, int(mc["input_length"]):, :, :, :1].squeeze(-1)
            loss = torch.mean(torch.abs(pred.squeeze(-1) - target))
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            optimizer.step()
            loss_sum += loss.detach() * frames.shape[0]
            sample_count += frames.shape[0]
        if distributed:
            torch.distributed.all_reduce(loss_sum)
            torch.distributed.all_reduce(sample_count)
        epoch_loss = (loss_sum / sample_count.clamp_min(1)).item()
        if rank == 0:
            print(f"epoch={epoch + 1}/{epochs} loss={epoch_loss:.6f}")
            if epoch_loss < best_loss:
                best_loss = epoch_loss
                best_epoch = epoch + 1
                best_state = clone_state_dict(model)
            loss_history.append({"epoch": epoch + 1, "loss": epoch_loss})
            loss_history_path.write_text(
                json.dumps(
                    {
                        "epochs": loss_history,
                        "best_epoch": best_epoch,
                        "best_loss": best_loss,
                    },
                    indent=2,
                )
                + "\n"
            )
    if rank == 0:
        if best_state is None:
            raise RuntimeError("Training finished without producing a model state")
        checkpoint_path = checkpoint_dir / f"{checkpoint_prefix}.pth"
        torch.save(
            {"state_dict": best_state, "epoch": best_epoch, "best_loss": best_loss},
            checkpoint_path,
        )
        print(f"checkpoint={checkpoint_path}")
        print(f"loss_history={loss_history_path}")
    if distributed:
        torch.distributed.barrier()
        torch.distributed.destroy_process_group()


if __name__ == "__main__":
    main()