| """Train PrecipitationSRCNN with optional torch DistributedDataParallel.""" |
|
|
| 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 |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
| from model.precipitationsrcnn import FORMAT_VERSION, MODEL_NAME, bilinear_input, build_model, processed_loss |
|
|
|
|
| def load_data(config): |
| data = np.load(ROOT / config["data"]["root"] / "daily_precipitation.npz") |
| expected = (216, 488) |
| if str(data["format_version"]) != FORMAT_VERSION or tuple(data["target_grid"]) != expected: |
| raise ValueError("data version or target grid mismatch") |
| target = data["target_precipitation"] |
| elevation = data["elevation"] |
| if target.ndim != 4 or target.shape[1:] != (1, *expected) or elevation.shape != (1, 1, *expected): |
| raise ValueError("invalid precipitation/elevation shape") |
| return data |
|
|
|
|
| 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")) |
| requested = config["runtime"]["device"] |
| use_cuda = requested == "cuda" and torch.cuda.is_available() |
| if distributed: |
| backend = "nccl" if use_cuda else "gloo" |
| torch.distributed.init_process_group(backend=backend) |
| rank = torch.distributed.get_rank() if distributed else 0 |
| world_size = torch.distributed.get_world_size() if distributed else 1 |
| device = torch.device(f"cuda:{local_rank}" if use_cuda else "cpu") |
| if device.type == "cuda": |
| torch.cuda.set_device(device) |
| data = load_data(config) |
| count = int(config["data"]["train_samples"]) |
| indices = list(range(rank, count, world_size)) |
| if not indices: |
| raise ValueError("each DDP rank requires at least one training sample") |
| model = build_model(config["model"]).to(device) |
| train_model = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None) if distributed else model |
| optimizer = torch.optim.Adam(train_model.parameters(), lr=float(config["training"]["learning_rate"])) |
| losses = [] |
| elevation = torch.from_numpy(data["elevation"]).to(device) |
| for _ in range(int(config["training"]["epochs"])): |
| for index in indices: |
| coarse = torch.from_numpy(data["coarse_precipitation"][index:index + 1]).to(device) |
| target = torch.from_numpy(data["target_precipitation"][index:index + 1]).to(device) |
| inputs = bilinear_input(coarse, elevation, tuple(config["data"]["target_grid"])) |
| optimizer.zero_grad(set_to_none=True) |
| prediction = train_model(inputs) |
| loss = processed_loss(prediction, target, config["training"]["loss"], |
| float(config["training"]["exponential_alpha"]), |
| float(config["training"]["exponential_weight"]), |
| float(config["training"]["quantile"])) |
| if not torch.isfinite(loss): |
| raise FloatingPointError("non-finite training loss") |
| loss.backward(); optimizer.step(); losses.append(float(loss.detach())) |
| if distributed: |
| gathered = [None] * world_size |
| torch.distributed.all_gather_object(gathered, losses) |
| losses = [value for part in gathered for value in part] |
| if rank == 0: |
| checkpoint_path = ROOT / config["paths"]["checkpoint"] |
| checkpoint_path.parent.mkdir(parents=True, exist_ok=True) |
| checkpoint = {"format_version": FORMAT_VERSION, "model_name": MODEL_NAME, |
| "model_config": config["model"], "paper_model": config["paper_model"], |
| "target_grid": config["data"]["target_grid"], "model": model.state_dict(), |
| "loss": config["training"]["loss"], "seed": seed} |
| torch.save(checkpoint, checkpoint_path) |
| metrics_path = ROOT / config["paths"]["training_metrics"] |
| metrics_path.parent.mkdir(parents=True, exist_ok=True) |
| metrics = {"format_version": FORMAT_VERSION, "epochs": int(config["training"]["epochs"]), |
| "world_size": world_size, "loss": config["training"]["loss"], |
| "loss_history": losses, "final_loss": losses[-1]} |
| metrics_path.write_text(json.dumps(metrics, indent=2) + "\n") |
| print(f"checkpoint={checkpoint_path.relative_to(ROOT)} final_loss={losses[-1]:.6f}") |
| if distributed: |
| torch.distributed.barrier(); torch.distributed.destroy_process_group() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
