| """Train, calibrate, and checkpoint WoFS elastic-net logistic models.""" |
|
|
| import argparse |
| 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, Subset |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
| from model.wofsstormcal import WoFSStormCal |
|
|
|
|
| class HazardDataset(Dataset): |
| def __init__(self, path, config): |
| self.data = np.load(path) |
| if str(self.data["format_version"]) != config["data"]["format_version"]: |
| raise ValueError("incompatible WoFS storm-object data format") |
| count = len(self.data["features"]) |
| if self.data["features"].shape != (count, 113): |
| raise ValueError("features must have shape [N,113]") |
| if self.data["targets"].shape != (count, 3): |
| raise ValueError("targets must have shape [N,3]") |
| if self.data["lead_group"].shape != (count,): |
| raise ValueError("lead_group must have shape [N]") |
| if not np.isfinite(self.data["features"]).all() or not np.isfinite(self.data["targets"]).all(): |
| raise ValueError("data contain NaN or Inf") |
|
|
| def __len__(self): |
| return len(self.data["features"]) |
|
|
| def __getitem__(self, index): |
| return (torch.from_numpy(self.data["features"][index]).float(), |
| torch.from_numpy(self.data["targets"][index]).float(), |
| torch.as_tensor(self.data["lead_group"][index], dtype=torch.long)) |
|
|
|
|
| def device_from_config(config, rank=0): |
| 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(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--resume", action="store_true", help="restore model and optimizer state before training") |
| args = parser.parse_args() |
| config = yaml.safe_load((ROOT / "conf/config.yaml").read_text()) |
| seed = int(config["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) |
| dataset = HazardDataset(ROOT / config["data"]["root"] / "train.npz", config) |
| split = int(len(dataset) * (1 - float(config["train"]["calibration_fraction"]))) |
| fit_set = Subset(dataset, range(split)) |
| sampler = DistributedSampler(fit_set, shuffle=True) if distributed else None |
| loader = DataLoader(fit_set, batch_size=int(config["train"]["batch_size"]), sampler=sampler, |
| shuffle=sampler is None, num_workers=int(config["train"]["num_workers"])) |
| model = WoFSStormCal(int(config["model"]["calibration_points"])).to(device) |
| features = dataset.data["features"][:split] |
| groups = dataset.data["lead_group"][:split] |
| means, scales = [], [] |
| for group in range(2): |
| group_features = features[groups == group] |
| means.append(group_features.mean(0)); scales.append(group_features.std(0).clip(1e-6)) |
| model.set_normalization(torch.from_numpy(np.stack(means)).to(device), torch.from_numpy(np.stack(scales)).to(device)) |
| wrapped = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None) if distributed else model |
| optimizer = torch.optim.Adam(wrapped.parameters(), lr=float(config["train"]["learning_rate"])) |
| checkpoint_path = ROOT / config["paths"]["checkpoint"] |
| start_epoch, history = 0, [] |
| if args.resume: |
| checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False) |
| model.load_state_dict(checkpoint["model"]) |
| optimizer.load_state_dict(checkpoint["optimizer"]) |
| start_epoch = int(checkpoint["epoch"]) |
| history = checkpoint.get("history", []) |
| for epoch in range(start_epoch, start_epoch + int(config["train"]["epochs"])): |
| if sampler is not None: |
| sampler.set_epoch(epoch) |
| total, steps = 0.0, 0 |
| for batch_features, targets, lead_group in loader: |
| batch_features, targets, lead_group = batch_features.to(device), targets.to(device), lead_group.to(device) |
| active_model = wrapped.module if distributed else wrapped |
| logits = wrapped(batch_features, lead_group, False) |
| logits = torch.logit(logits.clamp(1e-6, 1 - 1e-6)) |
| loss = active_model.elastic_net_loss(logits, targets, config["train"]["l1_strength"], config["train"]["l2_strength"]) |
| 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()); steps += 1 |
| if rank == 0: |
| history.append({"epoch": epoch + 1, "elastic_net_loss": total / max(steps, 1)}) |
| model = wrapped.module if distributed else wrapped |
| if rank == 0: |
| calibration_features = torch.from_numpy(dataset.data["features"][split:]).float().to(device) |
| calibration_targets = torch.from_numpy(dataset.data["targets"][split:]).float().to(device) |
| calibration_groups = torch.from_numpy(dataset.data["lead_group"][split:]).long().to(device) |
| model.fit_calibration(calibration_features, calibration_targets, calibration_groups) |
| checkpoint_path.parent.mkdir(parents=True, exist_ok=True) |
| payload = {"model": model.state_dict(), "optimizer": optimizer.state_dict(), "epoch": start_epoch + int(config["train"]["epochs"]), |
| "history": history, "format_version": config["data"]["format_version"], |
| "model_metadata": {"input_shape": ["N", 113], "output_shape": ["N", 3], |
| "hazards": model.hazards, "lead_groups": model.lead_groups, |
| "ensemble_members": 18, "grid_spacing_km": 3, |
| "forecast_window_minutes": 30, "forecast_interval_minutes": 5}} |
| torch.save(payload, checkpoint_path) |
| metrics = ROOT / config["paths"]["training_metrics"] |
| metrics.parent.mkdir(parents=True, exist_ok=True) |
| metrics.write_text(json.dumps({"history": history, "calibration_samples": len(dataset) - split}, indent=2) + "\n") |
| print(f"checkpoint={checkpoint_path.relative_to(ROOT)} output_shape=(N,3) epoch={payload['epoch']}") |
| if distributed: |
| torch.distributed.destroy_process_group() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
