| """Train UnetDif with single-process or torchrun DDP execution.""" |
|
|
| import json |
| import os |
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import yaml |
| from torch.nn.parallel import DistributedDataParallel as DDP |
| from torch.utils.data import DataLoader, Dataset, DistributedSampler |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
| from model.unetdif import UnetDif, loss_components |
|
|
|
|
| class RainfallDataset(Dataset): |
| def __init__(self, path: Path, config): |
| archive = np.load(path, allow_pickle=False) |
| expected = (int(config["time_steps"]), int(config["channels"]), int(config["height"]), int(config["width"])) |
| if archive["inputs"].shape[1:] != expected or archive["targets"].shape[1:] != (expected[0], expected[2], expected[3]): |
| raise ValueError(f"expected day shapes {expected} and {(expected[0], expected[2], expected[3])}") |
| if len(np.unique(archive["group_id"])) != len(archive["group_id"]): |
| raise ValueError("each rain day must belong to one split group") |
| if archive["lead_hours"].tolist() != config["lead_hours"]: |
| raise ValueError("NPZ lead_hours do not match the configured eight 3-hour steps") |
| self.inputs = archive["inputs"].astype(np.float32, copy=False).reshape(-1, *expected[1:]) |
| self.targets = archive["targets"].astype(np.float32, copy=False).reshape(-1, expected[2], expected[3]) |
|
|
| def __len__(self): |
| return len(self.inputs) |
|
|
| def __getitem__(self, index): |
| return torch.from_numpy(self.inputs[index]), torch.from_numpy(self.targets[index]) |
|
|
|
|
| def get_device(local_rank=0, local_world_size=1): |
| use_cuda = torch.cuda.is_available() and torch.cuda.device_count() >= local_world_size |
| return torch.device("cuda", local_rank) if use_cuda else torch.device("cpu") |
|
|
|
|
| def reduce_values(values, device, distributed): |
| tensor = torch.tensor(values, device=device) |
| if distributed: |
| torch.distributed.all_reduce(tensor) |
| return tensor.cpu().tolist() |
|
|
|
|
| def main(): |
| config = yaml.safe_load((ROOT / "conf/config.yaml").read_text()) |
| torch.manual_seed(int(config["seed"])) |
| distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1 |
| local_rank = int(os.environ.get("LOCAL_RANK", "0")) |
| local_world_size = int(os.environ.get("LOCAL_WORLD_SIZE", "1")) |
| device = get_device(local_rank, local_world_size) |
| if distributed: |
| torch.distributed.init_process_group("nccl" if device.type == "cuda" else "gloo") |
| rank = torch.distributed.get_rank() if distributed else 0 |
| dataset = RainfallDataset(ROOT / config["data"]["root"] / "train.npz", config["data"]) |
| if distributed and len(dataset) < torch.distributed.get_world_size(): |
| raise ValueError("DDP requires at least one training item per process") |
| sampler = DistributedSampler(dataset) if distributed else None |
| loader = DataLoader(dataset, batch_size=int(config["train"]["batch_size"]), shuffle=sampler is None, |
| sampler=sampler, num_workers=int(config["train"]["num_workers"])) |
| model = UnetDif(**config["model"]).to(device) |
| if distributed: |
| model = DDP(model, device_ids=[local_rank] if device.type == "cuda" else None) |
| optimizer = torch.optim.Adam(model.parameters(), lr=float(config["train"]["learning_rate"])) |
| names = ["dry_focal", "false_alarm_focal", "positive_mse", "negative_mse", "dry_mae", "all_mae"] |
| normalizers = {name: 1.0 for name in names} |
| history = [] |
| for epoch in range(int(config["train"]["epochs"])): |
| if sampler is not None: |
| sampler.set_epoch(epoch) |
| model.train(); sums = {name: 0.0 for name in names}; batches = 0 |
| for inputs, target in loader: |
| inputs, target = inputs.to(device), target.to(device) |
| heads = model(inputs) |
| raw = loss_components(heads, inputs[:, int(config["data"]["precipitation_channel"])], target, |
| float(config["data"]["rain_threshold_mm_3h"]), |
| float(config["loss"]["focal_alpha"]), float(config["loss"]["focal_gamma"])) |
| loss = sum(raw[name] * normalizers[name] for name in names) |
| optimizer.zero_grad(set_to_none=True); loss.backward(); optimizer.step() |
| for name in names: |
| sums[name] += float(raw[name].detach()) |
| batches += 1 |
| reduced = reduce_values([sums[name] for name in names] + [batches], device, distributed) |
| total_batches = max(reduced[-1], 1) |
| averages = {name: reduced[index] / total_batches for index, name in enumerate(names)} |
| if epoch == 0: |
| normalizers = {name: 1.0 / max(value, 1e-6) for name, value in averages.items()} |
| record = {"epoch": epoch + 1, "components": averages, "normalizers": normalizers} |
| history.append(record) |
| if rank == 0: |
| print(json.dumps(record)) |
| checkpoint = ROOT / config["paths"]["checkpoint"] |
| checkpoint.parent.mkdir(parents=True, exist_ok=True) |
| state = model.module.state_dict() if distributed else model.state_dict() |
| torch.save({"epoch": epoch + 1, "model": state, "optimizer": optimizer.state_dict(), |
| "model_config": config["model"], "loss_normalizers": normalizers, |
| "loss_components": names, |
| "format_version": config["data"]["format_version"]}, checkpoint) |
| if rank == 0: |
| metrics = ROOT / config["paths"]["training_metrics"] |
| metrics.parent.mkdir(parents=True, exist_ok=True) |
| metrics.write_text(json.dumps({"history": history}, indent=2) + "\n") |
| if distributed: |
| torch.distributed.destroy_process_group() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
