| """Train independent CRAI ensemble members, optionally under torchrun DDP.""" |
|
|
| from pathlib import Path |
| import argparse |
| import json |
| import os |
| import random |
| import sys |
| import numpy as np |
| import torch |
| from torch import distributed as dist |
| from torch.nn.parallel import DistributedDataParallel |
| from torch.utils.data import DataLoader, Dataset, DistributedSampler |
| import yaml |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
|
|
| from model.crai_climateextremes import CRAIClimateExtremes |
|
|
|
|
| class ClimateDataset(Dataset): |
| def __init__(self, archive, limit): |
| self.observed = torch.from_numpy(archive["observed"][:limit]) |
| self.valid = torch.from_numpy(archive["valid_mask"][:limit]) |
| self.target = torch.from_numpy(archive["target"][:limit]) |
| land = torch.from_numpy(archive["europe_mask"])[None, None] |
| self.missing = land * (1.0 - self.valid) |
|
|
| def __len__(self): |
| return len(self.target) |
|
|
| def __getitem__(self, index): |
| return torch.cat((self.observed[index], self.valid[index]), 0), self.target[index], self.missing[index] |
|
|
|
|
| def load_config(path): |
| with open(path, encoding="utf-8") as handle: |
| return yaml.safe_load(handle) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--config", type=Path, default=ROOT / "conf/config.yaml") |
| parser.add_argument("--paper-model", action="store_true") |
| args = parser.parse_args() |
| config_path = args.config if args.config.is_absolute() else ROOT / args.config |
| cfg = load_config(config_path) |
| use_paper = args.paper_model or cfg["paper_model"] |
| batch_size = cfg["paper_batch_size"] if use_paper else cfg["batch_size"] |
| iterations = cfg["paper_iterations"] if use_paper else cfg["max_iterations"] |
| members = cfg["paper_ensemble_members"] if use_paper else cfg["ensemble_members"] |
| rank, world = int(os.getenv("RANK", 0)), int(os.getenv("WORLD_SIZE", 1)) |
| if world > 1: |
| dist.init_process_group("nccl" if torch.cuda.is_available() else "gloo") |
| device = torch.device(f"cuda:{int(os.getenv('LOCAL_RANK', 0))}" if torch.cuda.is_available() else "cpu") |
| if device.type == "cuda": |
| torch.cuda.set_device(device) |
| archive = np.load(ROOT / cfg["data_path"]) |
| dataset = ClimateDataset(archive, cfg["num_samples"]) |
| checkpoint_path = ROOT / cfg["checkpoint_path"] |
| if rank == 0: |
| checkpoint_path.parent.mkdir(parents=True, exist_ok=True) |
| (ROOT / "result/training").mkdir(parents=True, exist_ok=True) |
| if world > 1: |
| dist.barrier() |
| records, member_states = [], [] |
| for member in range(members): |
| seed = int(cfg["seed"]) + member |
| random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) |
| sampler = DistributedSampler(dataset, shuffle=True, seed=seed) if world > 1 else None |
| loader = DataLoader(dataset, batch_size=batch_size, shuffle=sampler is None, sampler=sampler) |
| model = CRAIClimateExtremes(cfg["base_channels"]).to(device) |
| if world > 1: |
| model = DistributedDataParallel(model, device_ids=[device.index] if device.type == "cuda" else None) |
| optimizer = torch.optim.Adam(model.parameters(), lr=cfg["learning_rate"]) |
| step, losses = 0, [] |
| for epoch in range(cfg["epochs"] if not use_paper else 10**9): |
| if sampler is not None: |
| sampler.set_epoch(epoch) |
| for inputs, target, missing in loader: |
| inputs, target, missing = inputs.to(device), target.to(device), missing.to(device) |
| prediction = model(inputs) |
| loss = (torch.abs(prediction - target) * missing).sum() / missing.sum().clamp_min(1) |
| optimizer.zero_grad(); loss.backward(); optimizer.step() |
| losses.append(float(loss.detach())) |
| step += 1 |
| if step >= iterations: |
| break |
| if step >= iterations: |
| break |
| raw_model = model.module if isinstance(model, DistributedDataParallel) else model |
| if rank == 0: |
| member_states.append({key: value.detach().cpu() for key, value in raw_model.state_dict().items()}) |
| records.append({"member": member, "iterations": step, "final_missing_mae": losses[-1]}) |
| if rank == 0: |
| torch.save( |
| { |
| "format_version": "1.0", |
| "model_config": {"base_channels": cfg["base_channels"]}, |
| "model": member_states, |
| }, |
| checkpoint_path, |
| ) |
| payload = {"paper_model": bool(use_paper), "world_size": world, "members": records} |
| (ROOT / "result/training/metrics.json").write_text(json.dumps(payload, indent=2) + "\n") |
| print(json.dumps(payload)) |
| if world > 1: |
| dist.destroy_process_group() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
