File size: 6,688 Bytes
0fa8141 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | """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()
|