Publish Streamflow-LSTM reproduction

.gitattributes

@@ -1,35 +1,3 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

conf/config.yaml

@@ -0,0 +1,44 @@
+{
+  "seed": 2026,
+  "data": {
+    "format_version": "streamflow-lstm-v1",
+    "path": "data/streamflow_fake.npz",
+    "gauges": ["BNDN5", "ARWN8", "TCCC1", "CARO2", "ESSC2", "NFDC1", "LABW4", "CLNK1", "TRAC2", "NFSW4"],
+    "input_variables": 23,
+    "history_steps": 28,
+    "step_hours": 6,
+    "forecast_steps": 40,
+    "train_samples": 48,
+    "validation_samples": 12,
+    "forecast_cases": 12
+  },
+  "model": {
+    "hidden_size": 16,
+    "num_layers": 3,
+    "dropout": 0.1
+  },
+  "paper_model": {
+    "hidden_size": 50,
+    "num_layers": 3,
+    "activations": ["relu", "relu", "tanh"],
+    "ensemble_members_per_gauge": 100,
+    "best_members": 5,
+    "epochs": 10
+  },
+  "training": {
+    "ensemble_members_per_gauge": 2,
+    "best_members": 2,
+    "epochs": 10,
+    "batch_size": 16,
+    "learning_rate": 0.001,
+    "loss": "MSE"
+  },
+  "runtime": {"device": "auto", "num_threads": 2},
+  "paths": {
+    "checkpoint": "result/checkpoints/streamflow_lstm.pt",
+    "training_metrics": "result/training/metrics.json",
+    "predictions": "result/output/predictions.npz",
+    "evaluation_metrics": "result/evaluation/metrics.json",
+    "comparison": "result/evaluation/comparison.png"
+  }
+}

config.json

@@ -0,0 +1,52 @@
+{
+  "model_name": "Streamflow-LSTM",
+  "model_type": "gauge_specific_stacked_lstm",
+  "architectures": ["StreamflowLSTM"],
+  "framework": "PyTorch",
+  "torch_dtype": "float32",
+  "transformers_version": "4.44.0",
+  "domain": "hydrology",
+  "task": "six-hourly-streamflow-forecasting",
+  "license": "apache-2.0",
+  "paper": {
+    "title": "Using a long short-term memory (LSTM) neural network to boost river streamflow forecasts over the western United States",
+    "doi": "10.5194/hess-26-5449-2022"
+  },
+  "implementation": {
+    "entry_point": "model/streamflow_lstm.py",
+    "scope": "core-method reduced-ensemble engineering reproduction",
+    "train_script": "scripts/train.py",
+    "inference_script": "scripts/inference.py",
+    "evaluation_script": "scripts/result.py",
+    "synthetic_data_script": "scripts/fake_data.py"
+  },
+  "architecture": {
+    "input_shape": ["B", 28, 23],
+    "output_shape": ["B"],
+    "paper_hidden_sizes": [50, 50, 50],
+    "activations": ["ReLU", "ReLU", "tanh"],
+    "dense_outputs": 1
+  },
+  "data": {
+    "format_version": "streamflow-lstm-v1",
+    "gauges": 10,
+    "input_shape": ["B", 28, 23],
+    "forecast_steps": 40,
+    "step_hours": 6,
+    "unit": "m3 s-1",
+    "synthetic": true
+  },
+  "checkpoint": {
+    "path": "result/checkpoints/streamflow_lstm.pt",
+    "format": "single-file multi-gauge multi-member PyTorch checkpoint",
+    "official_weights": false
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/streamflow_lstm.py",
+    "scripts/fake_data.py",
+    "scripts/train.py",
+    "scripts/inference.py",
+    "scripts/result.py"
+  ]
+}

model/streamflow_lstm.py

@@ -0,0 +1,75 @@
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch import nn
+
+
+def load_config(path):
+    with open(path, encoding="utf-8") as handle:
+        return json.load(handle)
+
+
+def write_json(path, value):
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with open(path, "w", encoding="utf-8") as handle:
+        json.dump(value, handle, indent=2, sort_keys=True)
+
+
+class StreamflowLSTM(nn.Module):
+    """Paper stack with explicit inter-layer activation semantics."""
+
+    def __init__(self, input_size=23, hidden_size=50, dropout=0.1):
+        super().__init__()
+        self.layers = nn.ModuleList([
+            nn.LSTM(input_size, hidden_size, batch_first=True),
+            nn.LSTM(hidden_size, hidden_size, batch_first=True),
+            nn.LSTM(hidden_size, hidden_size, batch_first=True),
+        ])
+        self.dropout = nn.Dropout(dropout)
+        self.dense = nn.Linear(hidden_size, 1)
+
+    def forward(self, inputs):
+        if inputs.ndim != 3 or inputs.shape[-2:] != (28, 23):
+            raise ValueError(f"expected [B,28,23], got {tuple(inputs.shape)}")
+        values, _ = self.layers[0](inputs)
+        values = self.dropout(torch.relu(values))
+        values, _ = self.layers[1](values)
+        values = self.dropout(torch.relu(values))
+        values, _ = self.layers[2](values)
+        values = torch.tanh(values[:, -1])
+        return self.dense(values).squeeze(-1)
+
+
+def nse(prediction, target):
+    prediction, target = np.asarray(prediction), np.asarray(target)
+    denominator = np.sum((target - target.mean()) ** 2)
+    return float(1.0 - np.sum((prediction - target) ** 2) / max(denominator, 1e-12))
+
+
+def metrics(prediction, target):
+    prediction = np.asarray(prediction, dtype=np.float64)
+    target = np.asarray(target, dtype=np.float64)
+    correlation = float(np.corrcoef(prediction, target)[0, 1]) if prediction.size > 1 else 0.0
+    if not np.isfinite(correlation):
+        correlation = 0.0
+    alpha = float(prediction.std() / max(target.std(), 1e-12))
+    beta = float(prediction.mean() / max(target.mean(), 1e-12))
+    kge = float(1.0 - np.sqrt((correlation - 1) ** 2 + (alpha - 1) ** 2 + (beta - 1) ** 2))
+    return {
+        "kge": kge,
+        "kge_r": correlation,
+        "kge_alpha": alpha,
+        "kge_beta": beta,
+        "nse": nse(prediction, target),
+        "rmse_m3_s": float(np.sqrt(np.mean((prediction - target) ** 2))),
+    }
+
+
+def load_member(payload, device):
+    model = StreamflowLSTM(hidden_size=payload["hidden_size"], dropout=payload["dropout"]).to(device)
+    model.load_state_dict(payload["state_dict"])
+    model.eval()
+    return model, payload

scripts/fake_data.py

@@ -0,0 +1,69 @@
+#!/usr/bin/env python3
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+
+parser = argparse.ArgumentParser(description="Generate deterministic hydrometeorological smoke-test data")
+parser.add_argument("--config", default="conf/config.yaml")
+args = parser.parse_args()
+with open(args.config, encoding="utf-8") as handle:
+    config = json.load(handle)
+
+rng = np.random.default_rng(config["seed"])
+data = config["data"]
+gauges, train_n, val_n = len(data["gauges"]), data["train_samples"], data["validation_samples"]
+cases, leads, history, variables = data["forecast_cases"], data["forecast_steps"], 28, 23
+means = np.array([0.37, 8.24, 20.5, 3.28, 4.47, 49.0, 26.0, 0.14, 1.19, 8.55], dtype=np.float32)
+
+
+def forcing(gauge, count, phase_offset=0.0):
+    phase = rng.uniform(0, 2 * np.pi, (count, 1)) + phase_offset
+    t = np.arange(history, dtype=np.float32)[None] / 4.0
+    x = rng.normal(0, 0.35, (count, history, variables)).astype(np.float32)
+    temperature = 9 + 12 * np.sin(2 * np.pi * (t / 365 + phase / (2 * np.pi))) - gauge * 0.35
+    precipitation = rng.gamma(1.2, 0.8, (count, history)) * (0.5 + gauge / 18)
+    x[:, :, 2] = temperature
+    x[:, :, 5] = precipitation
+    x[:, :, 6] = np.maximum(0, precipitation * 0.55 + rng.normal(0, 0.1, precipitation.shape))
+    x[:, :, 10:14] = np.clip(0.45 + np.cumsum(precipitation[:, :, None] * 0.004, axis=1), 0, 1)
+    response = means[gauge] + means[gauge] * (0.08 * precipitation[:, -8:].mean(1) + 0.025 * x[:, -1, 6])
+    response += rng.normal(0, max(float(means[gauge]) * 0.03, 0.02), count)
+    raw = np.maximum(0, response * (1.12 - gauge * 0.012) + rng.normal(0, max(float(means[gauge]) * 0.05, 0.03), count))
+    x[:, :, 20], x[:, :, 21], x[:, :, 22] = raw[:, None], (0.92 * raw)[:, None], means[gauge]
+    return x, np.maximum(0, response).astype(np.float32)
+
+
+train_x = np.empty((gauges, train_n, history, variables), np.float32)
+train_y = np.empty((gauges, train_n), np.float32)
+val_x = np.empty((gauges, val_n, history, variables), np.float32)
+val_y = np.empty((gauges, val_n), np.float32)
+forecast_x = np.empty((gauges, cases, leads, history, variables), np.float32)
+forecast_y = np.empty((gauges, cases, leads), np.float32)
+persistence = np.empty_like(forecast_y)
+glofas = np.empty_like(forecast_y)
+for g in range(gauges):
+    train_x[g], train_y[g] = forcing(g, train_n)
+    val_x[g], val_y[g] = forcing(g, val_n, 0.3)
+    for case in range(cases):
+        base_x, base_y = forcing(g, 1, case / 7)
+        initial = float(base_y[0])
+        persistence[g, case] = initial
+        state = initial
+        for lead in range(leads):
+            window, target = forcing(g, 1, case / 7 + lead / 80)
+            # Forecast meteorology evolves, while observed flow is frozen at issue time.
+            window[:, :, 20] = initial
+            window[:, :, 21] *= 1.0 + 0.003 * lead
+            state = 0.75 * state + 0.25 * float(target[0])
+            forecast_x[g, case, lead] = window[0]
+            forecast_y[g, case, lead] = state
+            glofas[g, case, lead] = max(0, state * (1.10 - 0.002 * lead) + rng.normal(0, max(float(means[g]) * 0.06, 0.03)))
+
+path = Path(data["path"])
+path.parent.mkdir(parents=True, exist_ok=True)
+np.savez_compressed(path, train_x=train_x, train_y=train_y, val_x=val_x, val_y=val_y,
+                    forecast_x=forecast_x, forecast_y=forecast_y, persistence=persistence,
+                    glofas=glofas, gauges=np.array(data["gauges"]), lead_hours=np.arange(1, leads + 1) * 6)
+print(f"wrote {path}: train={train_x.shape}, forecast={forecast_x.shape}")

scripts/inference.py

@@ -0,0 +1,54 @@
+#!/usr/bin/env python3
+import argparse
+import json
+from pathlib import Path
+import sys
+
+import numpy as np
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from model.streamflow_lstm import load_member
+
+parser = argparse.ArgumentParser(description="Generate 40-step streamflow forecasts")
+parser.add_argument("--config", default="conf/config.yaml")
+parser.add_argument("--paper", action="store_true")
+args = parser.parse_args()
+with open(args.config, encoding="utf-8") as handle:
+    config = json.load(handle)
+device = torch.device("cuda" if config["runtime"]["device"] == "auto" and torch.cuda.is_available() else "cpu")
+best_count = config["paper_model" if args.paper else "training"]["best_members"]
+with np.load(config["data"]["path"]) as data:
+    forecast_x, target = data["forecast_x"], data["forecast_y"]
+    persistence, glofas = data["persistence"], data["glofas"]
+    gauges, lead_hours = data["gauges"].astype(str), data["lead_hours"]
+prediction = np.empty_like(target)
+selected = {}
+checkpoint_path = Path(config["paths"]["checkpoint"])
+checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+if checkpoint.get("format_version") != config["data"]["format_version"]:
+    raise ValueError(f"checkpoint format does not match {config['data']['format_version']}")
+if checkpoint.get("gauges") != gauges.tolist():
+    raise ValueError("checkpoint gauge order does not match input data")
+for gauge_index, gauge in enumerate(gauges):
+    payloads = [item for item in checkpoint["members"] if item["gauge"] == gauge]
+    payloads.sort(key=lambda item: item["validation_nse"], reverse=True)
+    chosen = payloads[:best_count]
+    if not chosen:
+        raise FileNotFoundError(f"no members for {gauge} in {checkpoint_path}; run scripts/train.py first")
+    x = forecast_x[gauge_index].reshape(-1, 28, 23)
+    member_predictions = []
+    for payload in chosen:
+        model, payload = load_member(payload, device)
+        normalized = torch.from_numpy((x - payload["x_mean"]) / payload["x_std"]).to(device)
+        with torch.no_grad():
+            values = model(normalized).cpu().numpy() * payload["y_std"] + payload["y_mean"]
+        member_predictions.append(np.maximum(0, values.reshape(target.shape[1:])))
+    prediction[gauge_index] = np.mean(member_predictions, axis=0)
+    selected[gauge] = np.array([item["member"] for item in chosen], dtype=np.int64)
+path = Path(config["paths"]["predictions"])
+path.parent.mkdir(parents=True, exist_ok=True)
+np.savez_compressed(path, prediction=prediction, target=target, persistence=persistence,
+                    glofas=glofas, gauges=gauges, lead_hours=lead_hours,
+                    selected_json=np.array(json.dumps({key: value.tolist() for key, value in selected.items()})))
+print(f"wrote {path}: {prediction.shape}")

scripts/result.py

@@ -0,0 +1,57 @@
+#!/usr/bin/env python3
+import argparse
+import json
+from pathlib import Path
+import sys
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from model.streamflow_lstm import metrics, write_json
+
+parser = argparse.ArgumentParser(description="Evaluate streamflow forecasts and baselines")
+parser.add_argument("--config", default="conf/config.yaml")
+args = parser.parse_args()
+with open(args.config, encoding="utf-8") as handle:
+    config = json.load(handle)
+with np.load(config["paths"]["predictions"]) as data:
+    methods = {name: data[name] for name in ("prediction", "persistence", "glofas")}
+    target, gauges, lead_hours = data["target"], data["gauges"].astype(str), data["lead_hours"]
+lead_days = {"2_day": 7, "5_day": 19, "8_day": 31}
+report = {"lead_definition": "zero-based indices 7/19/31 equal 48/120/192 hours", "gauges": {}}
+for gauge_index, gauge in enumerate(gauges):
+    report["gauges"][gauge] = {}
+    for method, values in methods.items():
+        report["gauges"][gauge][method] = {
+            label: metrics(values[gauge_index, :, index], target[gauge_index, :, index])
+            for label, index in lead_days.items()
+        }
+numbers = [value for gauge in report["gauges"].values() for method in gauge.values()
+           for lead in method.values() for value in lead.values()]
+if not np.all(np.isfinite(numbers)):
+    raise FloatingPointError("evaluation metrics are not finite")
+write_json(config["paths"]["evaluation_metrics"], report)
+
+colors = {"prediction": "#146c94", "persistence": "#d17a22", "glofas": "#6a8e3a"}
+figure, axes = plt.subplots(2, 1, figsize=(11, 8), constrained_layout=True)
+for method, values in methods.items():
+    rmse = [np.mean([metrics(values[g, :, lead], target[g, :, lead])["rmse_m3_s"] for g in range(len(gauges))])
+            for lead in range(len(lead_hours))]
+    axes[0].plot(lead_hours / 24, rmse, label=method, color=colors[method], linewidth=2)
+axes[0].set(title="Mean gauge RMSE across forecast lead", xlabel="Lead (days)", ylabel="RMSE (m3 s-1)")
+axes[0].grid(alpha=0.25)
+axes[0].legend()
+x = np.arange(len(gauges))
+width = 0.25
+for offset, (method, values) in enumerate(methods.items()):
+    kge = [metrics(values[g, :, 19], target[g, :, 19])["kge"] for g in range(len(gauges))]
+    axes[1].bar(x + (offset - 1) * width, kge, width, label=method, color=colors[method])
+axes[1].axhline(1 - np.sqrt(2), color="black", linestyle="--", linewidth=1, label="KGE skill threshold")
+axes[1].set(title="Gauge KGE at 5-day lead", xlabel="Gauge", ylabel="KGE", xticks=x, xticklabels=gauges)
+axes[1].legend(ncol=4, fontsize=8)
+comparison = Path(config["paths"]["comparison"])
+comparison.parent.mkdir(parents=True, exist_ok=True)
+figure.savefig(comparison, dpi=150)
+plt.close(figure)
+print(config["paths"]["evaluation_metrics"], comparison)

scripts/train.py

@@ -0,0 +1,99 @@
+#!/usr/bin/env python3
+import argparse
+import json
+import os
+from pathlib import Path
+import sys
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader, TensorDataset
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from model.streamflow_lstm import StreamflowLSTM, nse, write_json
+
+parser = argparse.ArgumentParser(description="Train independent ensembles for all stream gauges")
+parser.add_argument("--config", default="conf/config.yaml")
+parser.add_argument("--paper", action="store_true", help="Use 50 units, 100 members, best 5")
+parser.add_argument("--members", type=int, default=None)
+args = parser.parse_args()
+with open(args.config, encoding="utf-8") as handle:
+    config = json.load(handle)
+torch.set_num_threads(config["runtime"]["num_threads"])
+rank = int(os.environ.get("RANK", 0))
+world_size = int(os.environ.get("WORLD_SIZE", 1))
+local_rank = int(os.environ.get("LOCAL_RANK", 0))
+distributed = world_size > 1
+use_cuda = config["runtime"]["device"] == "auto" and torch.cuda.is_available()
+if distributed:
+    torch.distributed.init_process_group(backend="nccl" if use_cuda else "gloo")
+if use_cuda:
+    torch.cuda.set_device(local_rank)
+device = torch.device(f"cuda:{local_rank}" if use_cuda else "cpu")
+mode = config["paper_model"] if args.paper else {**config["training"], **config["model"]}
+members = args.members or mode["ensemble_members_per_gauge"]
+hidden, epochs = mode["hidden_size"], mode["epochs"]
+with np.load(config["data"]["path"]) as data:
+    train_x, train_y = data["train_x"], data["train_y"]
+    val_x, val_y, gauges = data["val_x"], data["val_y"], data["gauges"].astype(str)
+records, trained_members = {}, []
+for gauge_index, gauge in enumerate(gauges):
+    x_mean = train_x[gauge_index].mean((0, 1), keepdims=True)
+    x_std = train_x[gauge_index].std((0, 1), keepdims=True) + 1e-6
+    y_mean, y_std = float(train_y[gauge_index].mean()), float(train_y[gauge_index].std() + 1e-6)
+    x_train = torch.from_numpy((train_x[gauge_index] - x_mean) / x_std)
+    y_train = torch.from_numpy((train_y[gauge_index] - y_mean) / y_std)
+    x_val = torch.from_numpy((val_x[gauge_index] - x_mean) / x_std).to(device)
+    loader = DataLoader(TensorDataset(x_train, y_train), batch_size=config["training"]["batch_size"], shuffle=True)
+    scores = []
+    for member in range(members):
+        if (gauge_index * members + member) % world_size != rank:
+            continue
+        seed = config["seed"] + gauge_index * 1000 + member
+        torch.manual_seed(seed)
+        model = StreamflowLSTM(hidden_size=hidden, dropout=config["model"]["dropout"]).to(device)
+        optimizer = torch.optim.Adam(model.parameters(), lr=config["training"]["learning_rate"])
+        losses = []
+        model.train()
+        for _ in range(epochs):
+            for inputs, targets in loader:
+                optimizer.zero_grad(set_to_none=True)
+                loss = torch.mean((model(inputs.to(device)) - targets.to(device)) ** 2)
+                loss.backward()
+                optimizer.step()
+                losses.append(float(loss.detach()))
+        model.eval()
+        with torch.no_grad():
+            prediction = model(x_val).cpu().numpy() * y_std + y_mean
+        score = nse(prediction, val_y[gauge_index])
+        payload = {"state_dict": {key: value.detach().cpu() for key, value in model.state_dict().items()},
+                   "hidden_size": hidden, "dropout": config["model"]["dropout"],
+                   "x_mean": x_mean, "x_std": x_std, "y_mean": y_mean, "y_std": y_std,
+                   "gauge": gauge, "member": member, "validation_nse": score}
+        trained_members.append(payload)
+        scores.append({"member": member, "validation_nse": score, "final_mse": losses[-1]})
+        print(f"gauge={gauge} member={member} mse={losses[-1]:.6f} val_nse={score:.4f}")
+    records[gauge] = scores
+if distributed:
+    gathered = [None] * world_size if rank == 0 else None
+    torch.distributed.gather_object((trained_members, records), gathered, dst=0)
+    if rank == 0:
+        trained_members = [item for members_and_records in gathered for item in members_and_records[0]]
+        records = {gauge: [] for gauge in gauges}
+        for _, rank_records in gathered:
+            for gauge, values in rank_records.items():
+                records[gauge].extend(values)
+if rank == 0:
+    trained_members.sort(key=lambda item: (item["gauge"], item["member"]))
+    for scores in records.values():
+        scores.sort(key=lambda item: item["validation_nse"], reverse=True)
+    checkpoint = Path(config["paths"]["checkpoint"])
+    checkpoint.parent.mkdir(parents=True, exist_ok=True)
+    torch.save({"format_version": config["data"]["format_version"], "gauges": gauges.tolist(),
+                "members_per_gauge": members, "paper_mode": args.paper,
+                "members": trained_members}, checkpoint)
+    write_json(config["paths"]["training_metrics"], {"paper_mode": args.paper, "hidden_size": hidden,
+               "epochs": epochs, "members_per_gauge": members, "world_size": world_size, "gauges": records})
+    print(f"wrote {checkpoint}: {len(trained_members)} members")
+if distributed:
+    torch.distributed.destroy_process_group()