scripts/train.py

"""Train MassConservingCNN with optional torchrun DDP."""

import json
import os
import random
import sys
from pathlib import Path

import numpy as np
import torch
import yaml
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, Dataset, DistributedSampler


ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from model.massconservingcnn import MassConservingCNN


class MSWDataset(Dataset):
    def __init__(self, path, config):
        self.data = np.load(path)
        if str(self.data["format_version"]) != config["data"]["format_version"]:
            raise ValueError("incompatible data format version")
        count = len(self.data["inputs"])
        if self.data["inputs"].shape != (count, 4, 250):
            raise ValueError("inputs must have shape [B,4,250]")
        if self.data["targets"].shape != (count, 3, 250):
            raise ValueError("targets must have shape [B,3,250]")
        if self.data["inputs"].dtype != np.float32 or self.data["targets"].dtype != np.float32:
            raise TypeError("inputs and targets must be float32")
        if not np.isfinite(self.data["inputs"]).all() or not np.isfinite(self.data["targets"]).all():
            raise ValueError("data must be finite")
        if not np.isin(self.data["radar"], (0.0, 1.0)).all():
            raise ValueError("radar indicator must be binary")
        if (self.data["inputs"][:, 2] < 0).any() or (self.data["targets"][:, 2] < 0).any():
            raise ValueError("normalized rain must remain non-negative")

    def __len__(self):
        return len(self.data["inputs"])

    def __getitem__(self, index):
        return torch.from_numpy(self.data["inputs"][index]), torch.from_numpy(self.data["targets"][index])


def paper_j(prediction, target):
    return torch.sqrt(torch.mean((prediction - target) ** 2, dim=2) + 1e-12).mean(dim=1)


def mass_aware_loss(prediction, target, eta):
    base = paper_j(prediction, target)
    mass = eta / prediction.shape[2] * torch.abs(prediction[:, 1].sum(1) - target[:, 1].sum(1))
    return (base + mass).mean(), base.mean(), mass.mean()


def device_from_config(config, local_rank=0):
    requested = config["runtime"]["device"]
    if requested == "auto":
        return torch.device("cuda", local_rank) if torch.cuda.is_available() else torch.device("cpu")
    return torch.device(requested)


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    seed = int(config["seed"])
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    if distributed:
        torch.distributed.init_process_group("nccl" if torch.cuda.is_available() else "gloo")
    rank = torch.distributed.get_rank() if distributed else 0
    device = device_from_config(config, local_rank)
    if device.type == "cuda":
        torch.cuda.set_device(device); torch.cuda.manual_seed_all(seed)
    train_set = MSWDataset(ROOT / config["data"]["root"] / "train.npz", config)
    valid_set = MSWDataset(ROOT / config["data"]["root"] / "validation.npz", config)
    sampler = DistributedSampler(train_set, shuffle=True, seed=seed) if distributed else None
    loader = DataLoader(train_set, batch_size=int(config["train"]["batch_size"]),
                        shuffle=sampler is None, sampler=sampler,
                        num_workers=int(config["train"]["num_workers"]))
    valid_loader = DataLoader(valid_set, batch_size=int(config["train"]["batch_size"]), shuffle=False)
    model = MassConservingCNN(**config["model"]).to(device)
    wrapped = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None) if distributed else model
    bare = wrapped.module if distributed else wrapped
    optimizer = torch.optim.Adam(wrapped.parameters(), lr=float(config["train"]["learning_rate"]))
    history = []
    for epoch in range(int(config["train"]["epochs"])):
        if sampler is not None:
            sampler.set_epoch(epoch)
        wrapped.train(); total = 0.0; seen = 0
        for inputs, targets in loader:
            prediction = wrapped(inputs.to(device)); loss, _, _ = mass_aware_loss(prediction, targets.to(device), float(config["train"]["eta"]))
            optimizer.zero_grad(set_to_none=True); loss.backward(); optimizer.step()
            total += float(loss.detach()) * len(inputs); seen += len(inputs)
        totals = torch.tensor([total, seen], dtype=torch.float64, device=device)
        if distributed:
            torch.distributed.all_reduce(totals)
        wrapped.eval(); valid_total = valid_j = valid_mass = 0.0; valid_seen = 0
        if rank == 0:
            with torch.no_grad():
                for inputs, targets in valid_loader:
                    loss, base, mass = mass_aware_loss(bare(inputs.to(device)), targets.to(device), float(config["train"]["eta"]))
                    valid_total += float(loss) * len(inputs); valid_j += float(base) * len(inputs)
                    valid_mass += float(mass) * len(inputs); valid_seen += len(inputs)
            history.append({"epoch": epoch + 1, "train_loss": float(totals[0] / totals[1]),
                            "validation_loss": valid_total / valid_seen, "validation_J": valid_j / valid_seen,
                            "validation_mass_penalty": valid_mass / valid_seen})
    if rank == 0:
        checkpoint_path = ROOT / config["paths"]["checkpoint"]
        metrics_path = ROOT / config["paths"]["training_metrics"]
        checkpoint_path.parent.mkdir(parents=True, exist_ok=True); metrics_path.parent.mkdir(parents=True, exist_ok=True)
        model_state = bare.state_dict()
        torch.save({"model": model_state, "model_state_dict": model_state,
                    "optimizer_state_dict": optimizer.state_dict(),
                    "model_config": config["model"], "epoch": int(config["train"]["epochs"]),
                    "eta": float(config["train"]["eta"]), "format_version": config["data"]["format_version"],
                    "variable_order": ["u", "h", "r"], "normalization": "u,h: center/scale; r: scale only",
                    "climate_mean_uh": train_set.data["climate_mean_uh"],
                    "climate_std_uhr": train_set.data["climate_std_uhr"], "seed": seed}, checkpoint_path)
        metrics_path.write_text(json.dumps({"history": history}, indent=2) + "\n")
        print(f"checkpoint={checkpoint_path.relative_to(ROOT)} validation_loss={history[-1]['validation_loss']:.6f}")
    if distributed:
        torch.distributed.destroy_process_group()


if __name__ == "__main__":
    main()