#!/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()