File size: 5,600 Bytes
7a7efc9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | """Fit two valid-time folds without caching full-station threshold probabilities."""
import json
import os
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.improver_aifs import ImproverAIFS
from fake_data import generate_chunk
def fit_fold(config, meta, fold, rank):
model = ImproverAIFS(config["model"])
stations = int(config["data"]["station_count"])
chunk_size = int(config["train"]["station_chunk_size"])
# Bias uses every station but only station-centre values, so no full-grid cache exists.
for start in range(0, stations, chunk_size):
section = slice(start, min(start + chunk_size, stations))
centres, analyses, elevation = generate_chunk(
meta, fold, section, True, int(config["seed"]), include_patch=False
)
model.fit_bias_chunk(
torch.from_numpy(centres), torch.from_numpy(analyses),
torch.from_numpy(elevation), section,
)
sample_count = max(int(config["train"]["calibration_stations"]), int(config["train"]["blend_stations"]))
sample_stations = np.linspace(0, stations - 1, sample_count, dtype=np.int64)
patches, analyses, elevation = generate_chunk(
meta, fold, sample_stations, True, int(config["seed"]), include_patch=True
)
truth = torch.from_numpy(analyses)
thresholds = [torch.tensor(values, dtype=torch.float32) for values in config["model"]["thresholds"]]
source_expected = []
for source in range(3):
patch = torch.from_numpy(patches[:, :, :, source])
patch[:, :, 0] += (-0.0098 * torch.from_numpy(elevation)).view(1, 1, -1, 1, 1)
patch -= model.bias[source, :, :, sample_stations].unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
source_expected.append(patch[..., 1, 1])
probabilities = []
for variable, values in enumerate(thresholds):
width = model.fuzzy_widths[variable]
probability_patch = ((patch[:, :, variable].unsqueeze(2) - values.view(1, 1, -1, 1, 1, 1) + width) / (2 * width)).clamp(0, 1)
probabilities.append(model.neighborhood(model.recursive_filter(probability_patch)))
model.fit_reliability(probabilities, truth, thresholds, source)
del probabilities
blend_count = int(config["train"]["blend_stations"])
blend_input = torch.stack(source_expected, dim=3)[..., :blend_count]
blend_truth = truth[..., :blend_count]
optimizer = torch.optim.Adam([model.blend_logits], lr=float(config["train"]["learning_rate"]))
history = []
for epoch in range(int(config["train"]["epochs"])):
loss = torch.mean((model.blend_expected(blend_input) - blend_truth) ** 2)
if not torch.isfinite(loss):
raise ValueError("non-finite blend loss")
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
history.append({"epoch": epoch + 1, "blend_mse_loss": float(loss.detach())})
record = {
"fold": fold, "held_out_valid_date": str(meta["valid_dates"][fold]),
"history_days": 30, "bias_stations": stations,
"calibration_stations": int(config["train"]["calibration_stations"]),
"blend_stations": blend_count, "rank": rank, "history": history,
}
print(f"rank={rank} fold={fold} history_days=30 bias_stations=569 calibration_stations={sample_count} loss={history[-1]['blend_mse_loss']:.6f}")
return model.state_dict(), record
def main():
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
meta = np.load(ROOT / config["data"]["root"] / "protocol.npz", allow_pickle=True)
if str(meta["format_version"]) != config["data"]["format_version"] or meta["history_dates"].shape != (2, 30):
raise ValueError("protocol requires two valid dates and complete 30-day histories")
distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
if distributed:
torch.distributed.init_process_group("gloo")
rank = torch.distributed.get_rank() if distributed else 0
world_size = torch.distributed.get_world_size() if distributed else 1
local = []
for fold in range(2):
if fold % world_size == rank:
state, record = fit_fold(config, meta, fold, rank)
local.append((fold, state, record))
if distributed:
gathered = [None] * world_size
torch.distributed.all_gather_object(gathered, local)
combined = [item for rank_items in gathered for item in rank_items]
else:
combined = local
if rank == 0:
combined.sort(key=lambda item: item[0])
if [item[0] for item in combined] != [0, 1]:
raise RuntimeError("DDP ranks did not produce both valid-time folds")
fold_states = [item[1] for item in combined]
records = [item[2] for item in combined]
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": {"fold_states": fold_states, "fold_semantics": "valid-time 2-fold"},
"model_config": config["model"],
"format_version": config["data"]["format_version"],
}, checkpoint)
metrics.write_text(json.dumps({"folds": records, "distributed_world_size": world_size}, indent=2) + "\n")
if distributed:
torch.distributed.destroy_process_group()
if __name__ == "__main__":
main()
|