File size: 3,596 Bytes
e66b66f | 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 | """Run the paper-accurate no-training DataLoader diagnostic dry run."""
import json
import os
import sys
from pathlib import Path
import numpy as np
import torch
import yaml
from torch.utils.data import DataLoader, Dataset
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from model.pangu_icon_dke import PanguIconDKEDiagnostics
def load_protocol(config):
source = np.load(ROOT / config["data"]["root"] / config["data"]["protocol_file"])
protocol = {key: source[key].tolist() for key in source.files if key not in {"phase", "amplitude", "growth_rate", "spectral_phase", "spectral_amplitude", "base_wind_ms", "spectral_slope"}}
coefficients = {key: source[key].tolist() for key in ("phase", "amplitude", "growth_rate", "spectral_phase", "spectral_amplitude", "base_wind_ms", "spectral_slope")}
return protocol, coefficients
class ExperimentDataset(Dataset):
def __init__(self, count): self.count = count
def __len__(self): return self.count
def __getitem__(self, index): return torch.tensor(index, dtype=torch.long)
def main():
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
protocol, coefficients = load_protocol(config)
if tuple(protocol["field_shape"]) != (5, 73, 1, 721, 1440) or tuple(protocol["spectral_shape"]) != (5, 73, 1, 259560, 2):
raise ValueError("protocol dimensions do not match the public evaluation data")
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
dataset = ExperimentDataset(int(protocol["field_shape"][0]))
sampler = range(rank, len(dataset), torch.distributed.get_world_size()) if distributed else None
loader = DataLoader(dataset, batch_size=1, sampler=sampler, shuffle=False)
model = PanguIconDKEDiagnostics(protocol, coefficients)
if sum(parameter.numel() for parameter in model.parameters()) != 0:
raise RuntimeError("evaluation diagnostic must not have learnable parameters")
dry_values = {}
with torch.no_grad():
for item in loader:
experiment = int(item[0])
u, v = model.fields.wind_chunk(experiment, 0, slice(0, 8))
dry_values[experiment] = float(model(torch.from_numpy(u), torch.from_numpy(v)).mean())
if distributed:
gathered = [None] * torch.distributed.get_world_size() if rank == 0 else None
torch.distributed.gather_object(dry_values, gathered, dst=0)
if rank == 0:
dry_values = {key: value for shard in gathered for key, value in shard.items()}
if rank == 0:
checkpoint_path = ROOT / config["paths"]["checkpoint"]
metrics_path = ROOT / config["paths"]["training_metrics"]
checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
metrics_path.parent.mkdir(parents=True, exist_ok=True)
torch.save({"model": model.state_dict(), "model_config": config["generator"], "format_version": protocol["format_version"], "protocol": protocol, "coefficients": coefficients}, checkpoint_path)
metrics_path.write_text(json.dumps({"training_required": False, "backward_executed": False, "learnable_parameters": 0, "diagnostic_dry_run": True, "experiment_dry_run": [dry_values[index] for index in range(len(dataset))]}, indent=2) + "\n")
print(f"checkpoint={checkpoint_path.relative_to(ROOT)} training_required=false backward=false")
if distributed:
torch.distributed.destroy_process_group()
if __name__ == "__main__":
main()
|