| """Fit stable target clusters, NumPy RF, and neural classifiers.""" |
|
|
| import argparse |
| import json |
| import os |
| import random |
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import yaml |
| from torch import nn |
| from torch.nn.parallel import DistributedDataParallel |
| from torch.utils.data import DataLoader, TensorDataset |
| from torch.utils.data.distributed import DistributedSampler |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
| from model.cesm_seasonal_ml import (CHECKPOINT_FORMAT_VERSION, DATA_FORMAT_VERSION, FeedForwardNN, |
| NumpyGradientBoostedTrees, NumpyRandomForest, SeasonalLSTM, |
| StablePrecipKMeans, classification_metrics) |
|
|
|
|
| def setup(config): |
| distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1 |
| local_world = int(os.environ.get("LOCAL_WORLD_SIZE", "1")) |
| use_cuda = (torch.cuda.is_available() and config["runtime"]["device"] != "cpu" |
| and torch.cuda.device_count() >= local_world) |
| if distributed: |
| torch.distributed.init_process_group(config["runtime"]["ddp_backend_gpu"] if use_cuda else config["runtime"]["ddp_backend_cpu"]) |
| rank = torch.distributed.get_rank() if distributed else 0 |
| device = torch.device(f"cuda:{int(os.environ.get('LOCAL_RANK', 0))}" if use_cuda else "cpu") |
| if use_cuda: |
| torch.cuda.set_device(device) |
| return distributed, rank, device |
|
|
|
|
| def train_neural(name, model, x_train, y_train, x_val, y_val, settings, distributed, device): |
| dataset = TensorDataset(torch.from_numpy(x_train).float(), torch.from_numpy(y_train).long()) |
| sampler = DistributedSampler(dataset, shuffle=True) if distributed else None |
| local_samples = len(sampler) if sampler is not None else len(dataset) |
| loader = DataLoader(dataset, batch_size=settings["batch_size"], shuffle=sampler is None, |
| sampler=sampler, drop_last=name == "nn" and local_samples % settings["batch_size"] == 1) |
| model = model.to(device) |
| wrapped = DistributedDataParallel(model, device_ids=[device.index] if device.type == "cuda" else None) if distributed else model |
| counts = np.bincount(y_train, minlength=4) |
| weights = len(y_train) / np.maximum(4 * counts, 1) |
| criterion = nn.CrossEntropyLoss(weight=torch.tensor(weights, dtype=torch.float32, device=device)) |
| optimizer = torch.optim.Adam(wrapped.parameters(), lr=settings["learning_rate"]) |
| history, best_accuracy, best_state = [], -1.0, None |
| for epoch in range(settings["epochs"]): |
| if sampler: |
| sampler.set_epoch(epoch) |
| wrapped.train() |
| total_loss, total_samples = 0.0, 0 |
| for features, target in loader: |
| optimizer.zero_grad(set_to_none=True) |
| loss = criterion(wrapped(features.to(device)), target.to(device)) |
| loss.backward() |
| optimizer.step() |
| total_loss += float(loss.detach()) * len(features) |
| total_samples += len(features) |
| loss_summary = torch.tensor([total_loss, total_samples], dtype=torch.float64, device=device) |
| if distributed: |
| torch.distributed.all_reduce(loss_summary) |
| wrapped.eval() |
| with torch.no_grad(): |
| logits = wrapped(torch.from_numpy(x_val).float().to(device)) |
| correct = (logits.argmax(1).cpu().numpy() == y_val).sum() |
| count = len(y_val) |
| if distributed: |
| summary = torch.tensor([correct, count], dtype=torch.float64, device=device) |
| torch.distributed.all_reduce(summary) |
| correct, count = summary.tolist() |
| accuracy = correct / count |
| history.append({"epoch": epoch + 1, "loss": float(loss_summary[0] / loss_summary[1]), "validation_accuracy": accuracy}) |
| if accuracy > best_accuracy: |
| best_accuracy = accuracy |
| module = wrapped.module if distributed else wrapped |
| best_state = {key: value.detach().cpu().clone() for key, value in module.state_dict().items()} |
| return best_state, history |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--season", choices=["NDJ", "JFM"]) |
| parser.add_argument("--models", nargs="+", choices=["rf", "nn", "lstm", "xgboost"]) |
| args = parser.parse_args() |
| config = yaml.safe_load((ROOT / "conf/config.yaml").read_text()) |
| season = args.season or config["season"] |
| enabled = args.models or config["model"]["enabled"] |
| seed = int(config["seed"]) |
| random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) |
| distributed, rank, device = setup(config) |
| data = np.load(ROOT / config["data"]["path"]) |
| if str(data["format_version"]) != DATA_FORMAT_VERSION: |
| raise ValueError("data format version mismatch") |
| season_index = [str(x) for x in data["seasons"]].index(season) |
| train, validation, test = data["split"] == 0, data["split"] == 1, data["split"] == 2 |
| cluster = StablePrecipKMeans(seed=seed).fit(data["precipitation"][season_index, train]) |
| labels = cluster.transform(data["precipitation"][season_index]) |
| models, histories, validation_probabilities = {}, {}, {} |
| if "rf" in enabled: |
| rf_config = config["model"]["rf"] |
| forest = NumpyRandomForest(rf_config["trees"], rf_config["mtry"], rf_config["max_depth"], rf_config["min_samples_split"], seed) |
| if rank == 0: |
| forest.fit(data["rf_features"][season_index, train], labels[train]) |
| validation_probabilities["rf"] = forest.predict_proba(data["rf_features"][season_index, validation]) |
| importance = forest.permutation_importance(data["rf_features"][season_index, validation], labels[validation], rf_config["permutation_repeats"]) |
| mean_minimum_depth, root_frequency = forest.tree_structure_importance() |
| models["rf"] = forest |
| histories["rf"] = {"permutation_importance": importance.tolist(), "mean_minimum_depth": mean_minimum_depth.tolist(), |
| "root_frequency": root_frequency.tolist(), "feature_names": data["rf_manifest"].tolist()} |
| if "xgboost" in enabled and rank == 0: |
| xgb_config = config["model"]["xgboost"] |
| boosted = NumpyGradientBoostedTrees(xgb_config["n_rounds"], xgb_config["max_depth"], |
| xgb_config["eta"], xgb_config["gamma"], seed) |
| boosted.fit(data["rf_features"][season_index, train], labels[train]) |
| models["xgboost"] = boosted |
| validation_probabilities["xgboost"] = boosted.predict_proba(data["rf_features"][season_index, validation]) |
| histories["xgboost"] = {"rounds": xgb_config["n_rounds"]} |
| history_length = 4 if season == "NDJ" else 12 |
| specifications = { |
| "nn": (FeedForwardNN(hidden=tuple(config["model"]["nn"]["hidden"]), dropout=config["model"]["nn"]["dropout"]), data["nn_features"][season_index], config["model"]["nn"]), |
| "lstm": (SeasonalLSTM(hidden_size=config["model"]["lstm"]["hidden_size"], dense_size=config["model"]["lstm"]["dense_size"], dropout=config["model"]["lstm"]["dropout"]), data["eof_sequence"][season_index, :, -history_length:], config["model"]["lstm"]) |
| } |
| for name in ("nn", "lstm"): |
| if name in enabled: |
| model, features, settings = specifications[name] |
| state, history = train_neural(name, model, features[train], labels[train], features[validation], labels[validation], settings, distributed, device) |
| model.load_state_dict(state); model.to(device).eval() |
| with torch.no_grad(): |
| probability = torch.softmax(model(torch.from_numpy(features[validation]).float().to(device)), dim=1).cpu().numpy() |
| models[name] = state |
| histories[name] = history |
| validation_probabilities[name] = probability |
| if rank == 0: |
| metrics = {name: classification_metrics(labels[validation], probability, config["evaluation"]["random_trials"], seed) for name, probability in validation_probabilities.items()} |
| model_config = {"enabled": list(enabled), "primary": config["model"]["primary"], "settings": config["model"], |
| "history_length": history_length, "classes": 4, "rf_features": 103, |
| "nn_features": 416, "eof_channels": 28} |
| data_spec = {"format_version": str(data["format_version"]), "season": season, |
| "seasons": data["seasons"].tolist(), "rf_manifest": data["rf_manifest"].tolist(), |
| "nn_manifest": data["nn_manifest"].tolist(), "rf_shape": list(data["rf_features"].shape), |
| "nn_shape": list(data["nn_features"].shape), "eof_shape": list(data["eof_sequence"].shape), |
| "precipitation_shape": list(data["precipitation"].shape), "class_names": data["class_names"].tolist()} |
| checkpoint = {"format_version": CHECKPOINT_FORMAT_VERSION, "model": models, "model_config": model_config, |
| "data_spec": data_spec, "season": season, "cluster": cluster.state_dict(), |
| "training_history": histories, "validation_metrics": metrics, "paper_model": config["paper_model"]} |
| path = ROOT / config["paths"]["checkpoint"] |
| path.parent.mkdir(parents=True, exist_ok=True) |
| torch.save(checkpoint, path) |
| metrics_path = ROOT / config["paths"]["training_metrics"] |
| metrics_path.parent.mkdir(parents=True, exist_ok=True) |
| metrics_path.write_text(json.dumps({"season": season, "models": enabled, "validation": metrics}, indent=2) + "\n") |
| print(f"checkpoint={path.relative_to(ROOT)} season={season} models={','.join(enabled)}") |
| if distributed: |
| torch.distributed.barrier(); torch.distributed.destroy_process_group() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
