| """Train all 21 LOSO folds; DDP ranks independently shard outer folds.""" |
|
|
| import json |
| import argparse |
| import os |
| import random |
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import yaml |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
| from model.climemu_s2l import (FEATURE_COUNT, FORMAT_VERSION, GRID_SHAPE, |
| MODEL_NAME, DualRidge, SharedKernelGPR, |
| area_weights, select_ridge_alpha, weighted_rmse) |
|
|
|
|
| def load_data(path): |
| data = np.load(path) |
| short, long = data["short_response"], data["long_response"] |
| if str(data["format_version"]) != FORMAT_VERSION: |
| raise ValueError("data format version mismatch") |
| if short.shape != (21, *GRID_SHAPE) or long.shape != short.shape: |
| raise ValueError("responses must both have shape [21,145,192]") |
| ids = [str(value) for value in data["scenario_ids"]] |
| if len(set(ids)) != 21 or not np.isfinite(short).all() or not np.isfinite(long).all(): |
| raise ValueError("scenario IDs must be unique and fields finite") |
| return data, short.reshape(21, FEATURE_COUNT), long.reshape(21, FEATURE_COUNT), ids |
|
|
|
|
| def train_fold(fold, x, y, scenario_ids, config, weights): |
| training = np.asarray([index for index in range(21) if index != fold]) |
| alpha, cv_scores = select_ridge_alpha( |
| x[training], y[training], config["model"]["ridge"]["alphas"], |
| int(config["model"]["ridge"]["inner_folds"]), weights) |
| ridge = DualRidge(alpha).fit(x[training], y[training]) |
| ridge_prediction = ridge.predict(x[fold:fold + 1]).detach().cpu().numpy()[0] |
| gpr_config = config["model"]["gpr"] |
| gpr = SharedKernelGPR(gpr_config["kernel_mode"], gpr_config["jitter"]) |
| optimizer_trace = gpr.fit(x[training], y[training], gpr_config["optimizer_steps"], gpr_config["learning_rate"]) |
| gpr_prediction = gpr.predict(x[fold:fold + 1]).detach().cpu().numpy()[0] |
| state = {"fold": fold, "held_out_scenario_id": scenario_ids[fold], "train_indices": training.tolist(), |
| "ridge_alpha": alpha, "ridge_cv_scores": cv_scores, |
| "gpr_kernel_mode": gpr.kernel_mode, "gpr_hyperparameters": gpr.hyperparameters(), |
| "gpr_optimizer": optimizer_trace} |
| metrics = {"scenario_id": scenario_ids[fold], |
| "ridge_weighted_rmse": weighted_rmse(y[fold], ridge_prediction, weights), |
| "gpr_weighted_rmse": weighted_rmse(y[fold], gpr_prediction, weights), |
| "selected_alpha": alpha, "gpr_final_loss": optimizer_trace["loss"][-1], |
| "gpr_backward": optimizer_trace["gradient_seen"], |
| "gpr_parameter_update": optimizer_trace["parameter_updated"]} |
| return state, metrics |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--resume", action="store_true", help="resume completed LOSO folds from the checkpoint") |
| args = parser.parse_args() |
| 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) |
| torch.set_default_dtype(torch.float64) |
| distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1 |
| if distributed: |
| torch.distributed.init_process_group(config["runtime"]["ddp_backend"]) |
| rank = torch.distributed.get_rank() if distributed else 0 |
| world = torch.distributed.get_world_size() if distributed else 1 |
| data, x_numpy, y_numpy, scenario_ids = load_data(ROOT / config["data"]["path"]) |
| x, y = torch.from_numpy(x_numpy).double(), torch.from_numpy(y_numpy).double() |
| weights = area_weights(data["latitude_deg"], GRID_SHAPE[1]) |
| checkpoint_path = ROOT / config["paths"]["checkpoint"] |
| resumed_states, resumed_metrics = {}, {} |
| if args.resume and checkpoint_path.exists(): |
| try: |
| previous = torch.load(checkpoint_path, map_location="cpu", weights_only=False) |
| except TypeError: |
| previous = torch.load(checkpoint_path, map_location="cpu") |
| if previous.get("format_version") != FORMAT_VERSION or tuple(previous.get("grid_shape", ())) != GRID_SHAPE: |
| raise ValueError("checkpoint protocol mismatch") |
| restored_folds = previous.get("model", {}).get("folds", previous.get("folds", [])) |
| resumed_states = {state["fold"]: state for state in restored_folds} |
| resumed_metrics = {index: value for index, value in enumerate(previous.get("training_metrics", []))} |
| local_states, local_metrics = {}, {} |
| for fold in range(rank, 21, world): |
| if fold in resumed_states and fold in resumed_metrics: |
| continue |
| state, metrics = train_fold(fold, x, y, scenario_ids, config, weights) |
| local_states[fold], local_metrics[fold] = state, metrics |
| print(f"rank={rank} fold={fold:02d} scenario={scenario_ids[fold]} alpha={state['ridge_alpha']}") |
| if distributed: |
| gathered_states, gathered_metrics = [None] * world, [None] * world |
| torch.distributed.all_gather_object(gathered_states, local_states) |
| torch.distributed.all_gather_object(gathered_metrics, local_metrics) |
| states = {**resumed_states, **{key: value for item in gathered_states for key, value in item.items()}} |
| metrics = {**resumed_metrics, **{key: value for item in gathered_metrics for key, value in item.items()}} |
| else: |
| states, metrics = {**resumed_states, **local_states}, {**resumed_metrics, **local_metrics} |
| if rank == 0: |
| if sorted(states) != list(range(21)): |
| raise RuntimeError("all 21 LOSO folds must be trained") |
| fold_states = [states[index] for index in range(21)] |
| fold_metrics = [metrics[index] for index in range(21)] |
| checkpoint = {"epoch": 1, |
| "model": {"folds": fold_states}, |
| "model_config": config["model"], |
| "optimizer_state_dict": {"folds": [state["gpr_optimizer"] for state in fold_states]}, |
| "loss": float(np.mean([item["gpr_final_loss"] for item in fold_metrics])), |
| "config": config, |
| "model_name": MODEL_NAME, "format_version": FORMAT_VERSION, |
| "grid_shape": GRID_SHAPE, "scenario_ids": scenario_ids, |
| "short_response": torch.from_numpy(x_numpy), "long_response": torch.from_numpy(y_numpy), |
| "training_metrics": fold_metrics, |
| "gpr_ard": False, "gpr_ard_gap": config["model"]["gpr"]["ard_gap"]} |
| checkpoint_path.parent.mkdir(parents=True, exist_ok=True) |
| torch.save(checkpoint, checkpoint_path) |
| metrics_path = ROOT / config["paths"]["training_metrics"] |
| metrics_path.parent.mkdir(parents=True, exist_ok=True) |
| metrics_path.write_text(json.dumps({"fold_count": 21, "world_size": world, |
| "folds": fold_metrics}, indent=2) + "\n") |
| print(f"checkpoint={checkpoint_path.relative_to(ROOT)} folds=21 world_size={world} resumed={len(resumed_states)}") |
| if distributed: |
| torch.distributed.barrier(); torch.distributed.destroy_process_group() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
