| """Fit the paper's six station-wise model configurations.""" |
|
|
| 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.globalsurgeml import GlobalSurgeML |
|
|
|
|
| CONFIGURATIONS = { |
| "LR-RS": ("rs_daily", "linear"), |
| "LR-RS-lag": ("rs_lagged", "linear"), |
| "RF-RS-lag": ("rs_lagged", "random_forest"), |
| "LR-AR": ("ar_daily", "linear"), |
| "LR-AR-lag": ("ar_lagged", "linear"), |
| "RF-AR-lag": ("ar_lagged", "random_forest"), |
| } |
|
|
|
|
| class SurgeDataset(Dataset): |
| def __init__(self, path, config): |
| self.data = np.load(path) |
| expected = config["data"] |
| if str(self.data["format_version"]) != expected["format_version"]: |
| raise ValueError("incompatible storm-surge data format") |
| dimensions = {"rs_daily": expected["rs_daily_features"], "rs_lagged": expected["rs_lagged_features"], |
| "ar_daily": expected["ar_daily_features"], "ar_lagged": expected["ar_lagged_features"]} |
| for key, width in dimensions.items(): |
| if self.data[key].shape[1:] != (int(width),): |
| raise ValueError(f"{key} must have shape [N,{width}]") |
| if self.data["targets_m"].shape[1:] != (1,): |
| raise ValueError("targets_m must have shape [N,1]") |
|
|
| def __len__(self): |
| return len(self.data["targets_m"]) |
|
|
| def __getitem__(self, index): |
| return {key: torch.from_numpy(self.data[key][index]).float() for key in |
| ("rs_daily", "rs_lagged", "ar_daily", "ar_lagged", "targets_m")} |
|
|
|
|
| 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(): |
| 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 = SurgeDataset(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"])) |
| full = dataset.data |
| means, scales, states, history = {}, {}, {}, [] |
| models = {} |
| for model_index, (name, (feature_key, method)) in enumerate(CONFIGURATIONS.items()): |
| array = full[feature_key].astype(np.float32) |
| means[name] = array.mean(0).astype(np.float32) |
| scales[name] = array.std(0).clip(1e-6).astype(np.float32) |
| standardized = (array - means[name]) / scales[name] |
| model = GlobalSurgeML(array.shape[1], method, config["model"], seed + model_index).to(device) |
| if method == "linear": |
| model.regressor.select_features(standardized, full["targets_m"][:, 0], |
| float(config["model"]["p_value_threshold"])) |
| 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"]), |
| weight_decay=float(config["train"]["weight_decay"])) |
| for epoch in range(int(config["train"]["epochs"])): |
| if sampler: |
| sampler.set_epoch(epoch) |
| total, steps = 0.0, 0 |
| for batch in loader: |
| features = (batch[feature_key].to(device) - torch.from_numpy(means[name]).to(device)) / torch.from_numpy(scales[name]).to(device) |
| target = batch["targets_m"].to(device) |
| prediction = wrapped(features) |
| loss = torch.nn.functional.mse_loss(prediction, target) |
| 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 |
| history.append({"model": name, "epoch": epoch + 1, "mse_m2": total / max(steps, 1)}) |
| model = wrapped.module if distributed else wrapped |
| else: |
| model.regressor.fit(standardized, full["targets_m"][:, 0]) |
| prediction = model(torch.from_numpy(standardized).to(device)) |
| history.append({"model": name, "epoch": 1, |
| "mse_m2": float(torch.nn.functional.mse_loss(prediction.cpu(), torch.from_numpy(full["targets_m"])).item())}) |
| models[name] = model |
| states[name] = model.state_dict() |
| if rank == 0: |
| checkpoint = ROOT / config["paths"]["checkpoint"] |
| metrics = ROOT / config["paths"]["training_metrics"] |
| checkpoint.parent.mkdir(parents=True, exist_ok=True) |
| metrics.parent.mkdir(parents=True, exist_ok=True) |
| torch.save({"model": states, "model_config": config["model"], "configurations": CONFIGURATIONS, |
| "feature_means": means, "feature_scales": scales, |
| "format_version": config["data"]["format_version"], "target_unit": "m"}, checkpoint) |
| metrics.write_text(json.dumps({"history": history}, indent=2) + "\n") |
| print(f"checkpoint={checkpoint.relative_to(ROOT)} models={len(states)}") |
| if distributed: |
| torch.distributed.destroy_process_group() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
