File size: 4,748 Bytes
7f71cfd | 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 | import json
import os
from pathlib import Path
import sys
import numpy as np
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, DistributedSampler, TensorDataset
import yaml
ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
sys.path.insert(0, str(ROOT))
from model.stablenn_phys import StableNNPhys, rollout, rollout_loss
CHECKPOINT_FORMAT_VERSION = "stablenn_phys_checkpoint_v1"
def main():
cfg = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
torch.manual_seed(cfg["seed"])
distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
world_size = int(os.environ.get("WORLD_SIZE", "1"))
use_accelerator = torch.cuda.is_available() and torch.cuda.device_count() >= world_size
if distributed:
dist.init_process_group("nccl" if use_accelerator else "gloo")
rank = dist.get_rank() if distributed else 0
local_rank = int(os.environ.get("LOCAL_RANK", "0"))
device = torch.device(f"cuda:{local_rank}" if use_accelerator else "cpu")
if use_accelerator:
torch.cuda.set_device(device)
raw = np.load(ROOT / cfg["data"]["file"])
tensors = [torch.from_numpy(raw[name].astype(np.float32)) for name in ("initial", "target", "surface", "advection")]
source = torch.from_numpy(raw["source"].astype(np.int64))
dataset = TensorDataset(*tensors, source)
sampler = DistributedSampler(dataset, shuffle=True, seed=cfg["seed"]) if distributed else None
loader = DataLoader(dataset, batch_size=cfg["train"]["batch_size"], sampler=sampler,
shuffle=sampler is None)
all_state = torch.from_numpy(raw["target"].astype(np.float32))
state_mean = all_state.mean((0, 1)).to(device)
state_std = all_state.std((0, 1), unbiased=False).clamp_min(1e-5).to(device)
tendency_mean = torch.zeros(68, device=device)
tendency_std = torch.cat((torch.full((34,), 0.02), torch.full((34,), 2e-7))).to(device)
layer_mass = torch.from_numpy(raw["layer_mass"].astype(np.float32)).to(device)
model = StableNNPhys(cfg["model"]["hidden_size"]).to(device)
if distributed:
model = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None)
optimizer = torch.optim.Adam(model.parameters(), lr=cfg["train"]["learning_rate"])
history = []
for epoch in range(cfg["train"]["epochs"]):
if sampler:
sampler.set_epoch(epoch)
total = 0.0
for initial, target, surface, advection, _ in loader:
initial, target, surface, advection = [x.to(device) for x in (initial, target, surface, advection)]
prediction, _ = rollout(model, initial, surface, advection, state_mean, state_std,
tendency_mean, tendency_std, float(raw["dt_seconds"]))
loss = rollout_loss(prediction, target, layer_mass.expand(initial.shape[0], -1), cfg["train"]["loss_mode"])
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
total += loss.item() * initial.shape[0]
value = total / len(dataset)
history.append(value)
if rank == 0:
print(f"epoch={epoch + 1} loss={value:.6g}")
if rank == 0:
module = model.module if distributed else model
checkpoint = ROOT / cfg["paths"]["checkpoint"]
checkpoint.parent.mkdir(parents=True, exist_ok=True)
payload = {"format_version": CHECKPOINT_FORMAT_VERSION, "model": module.state_dict(),
"optimizer_state_dict": optimizer.state_dict(), "epoch": cfg["train"]["epochs"],
"model_config": dict(cfg["model"]), "training_config": dict(cfg["train"]),
"normalization": {"state_mean": state_mean.cpu(), "state_std": state_std.cpu(),
"tendency_mean": tendency_mean.cpu(), "tendency_std": tendency_std.cpu()},
"variables": {"input": ["sL[34]", "qT[34]", "SHF", "LHF", "SOLIN"],
"output": ["dsL_dt[34]", "dqT_dt[34]"]}, "dt_seconds": float(raw["dt_seconds"])}
torch.save(payload, checkpoint)
metrics = ROOT / cfg["paths"]["training_metrics"]
metrics.parent.mkdir(parents=True, exist_ok=True)
metrics.write_text(json.dumps({"loss": history, "world_size": dist.get_world_size() if distributed else 1,
"paper_model": cfg["train"]["paper_model"]}, indent=2))
print(f"saved {checkpoint}")
if distributed:
dist.destroy_process_group()
if __name__ == "__main__":
main()
|