| """Train four independent paper CNN-LSTM branches with RMSprop and MSE.""" |
|
|
| import json |
| import os |
| 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.climatebench import ClimateBench, ClimateBenchBranch, TARGETS |
|
|
|
|
| class ClimateDataset(Dataset): |
| def __init__(self, path: Path, config: dict): |
| self.data = np.load(path) |
| expected = config["data"] |
| shape = (int(expected["time_steps"]), 4, int(expected["height"]), int(expected["width"])) |
| target_shape = (4, int(expected["height"]), int(expected["width"])) |
| if str(self.data["format_version"]) != expected["format_version"] or str(self.data["storage_layout"]) != "NTCHW": |
| raise ValueError("incompatible ClimateBench NPZ format or storage layout") |
| if self.data["inputs"].shape[1:] != shape or self.data["targets"].shape[1:] != target_shape: |
| raise ValueError(f"expected inputs [N,{shape}] and targets [N,{target_shape}]") |
| if self.data["inputs"].dtype != np.float32 or self.data["targets"].dtype != np.float32: |
| raise TypeError("inputs and targets must be float32") |
| if tuple(self.data["channel_names"].tolist()) != tuple(expected["channels"]): |
| raise ValueError("forcing channels must be [co2_cumulative,ch4,so2,bc]") |
|
|
| def __len__(self) -> int: |
| return len(self.data["inputs"]) |
|
|
| def __getitem__(self, index: int): |
| return torch.from_numpy(self.data["inputs"][index]), torch.from_numpy(self.data["targets"][index]) |
|
|
|
|
| def device_from_config(config: dict, rank: int = 0) -> torch.device: |
| requested = config["runtime"]["device"] |
| if requested == "auto": |
| return torch.device("cuda", rank) if torch.cuda.is_available() else torch.device("cpu") |
| return torch.device(requested) |
|
|
|
|
| def main() -> None: |
| config = yaml.safe_load((ROOT / "conf/config.yaml").read_text()) |
| np.random.seed(int(config["seed"])) |
| torch.manual_seed(int(config["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) |
| dataset = ClimateDataset(ROOT / config["data"]["root"] / "train.npz", config) |
| sampler = DistributedSampler(dataset, shuffle=True) if distributed else None |
| loader = DataLoader(dataset, batch_size=int(config["train"]["batch_size"]), sampler=sampler, |
| shuffle=sampler is None, num_workers=int(config["train"]["num_workers"])) |
| model = ClimateBench(int(config["data"]["height"]), int(config["data"]["width"])).to(device) |
| counts = model.parameter_counts() |
| expected_count = int(config["model"]["parameters_per_target"]) |
| if any(count != expected_count for count in counts.values()): |
| raise RuntimeError(f"parameter count mismatch: {counts}") |
| wrapped = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None) if distributed else model |
| optimizer = torch.optim.RMSprop(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) |
| total, batches = 0.0, 0 |
| for inputs, targets in loader: |
| prediction = wrapped(inputs.to(device)) |
| loss = torch.nn.functional.mse_loss(prediction, targets.to(device)) |
| optimizer.zero_grad(set_to_none=True) |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(wrapped.parameters(), float(config["train"]["gradient_clip_norm"])) |
| optimizer.step() |
| total += float(loss.detach()) |
| batches += 1 |
| history.append({"epoch": epoch + 1, "mse": total / max(batches, 1), "batches": batches}) |
| model = wrapped.module if distributed else wrapped |
| 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) |
| torch.save({"model": model.state_dict(), "parameter_counts": counts, "targets": TARGETS, |
| "format_version": config["data"]["format_version"], "storage_layout": "NTCHW", |
| "model_config": {"height": model.height, "width": model.width}}, checkpoint_path) |
| metrics_path.write_text(json.dumps({"history": history, "parameter_counts": counts}, indent=2) + "\n") |
| print(f"checkpoint={checkpoint_path.relative_to(ROOT)} batches={history[-1]['batches']} parameters={counts}") |
| if distributed: |
| torch.distributed.destroy_process_group() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
