File size: 3,380 Bytes
5b3329b 986404c 0face05 5b3329b 0face05 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c | 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 | """Generate a checkpoint-backed CorrDiff ensemble in device-sized batches."""
import argparse
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.corrdiff import CorrDiff
def scalar(archive, key):
if key not in archive or archive[key].ndim != 0:
raise ValueError(f"NPZ metadata {key} must be present as a scalar")
return str(archive[key].item())
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--config", default=str(ROOT / "conf/config.yaml"))
parser.add_argument("--checkpoint")
args = parser.parse_args()
config = yaml.safe_load(Path(args.config).read_text(encoding="utf-8"))
checkpoint_path = Path(args.checkpoint) if args.checkpoint else ROOT / config["paths"]["checkpoint"]
if not checkpoint_path.is_file():
raise FileNotFoundError(f"Checkpoint is required: {checkpoint_path}")
device = torch.device("cuda" if torch.cuda.is_available() and config["runtime"]["device"] != "cpu" else "cpu")
state = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
model = CorrDiff(**config["model"]).to(device)
model.load_state_dict(state["model"])
model.eval()
archive = np.load(ROOT / config["data"]["path"])
protocol, data_source = scalar(archive, "protocol"), scalar(archive, "data_source")
if protocol != config["data"]["protocol"] or state.get("protocol") != protocol:
raise ValueError("Data, checkpoint, and configured protocols must match")
if not data_source or state.get("data_source") != data_source:
raise ValueError("Data and checkpoint data_source metadata must match")
coarse = archive[config["data"]["input_key"]]
target = archive[config["data"]["target_key"]]
if coarse.ndim != 4 or target.ndim != 4 or tuple(coarse.shape[1:]) != tuple(config["data"]["input_shape"]) or tuple(target.shape[1:]) != tuple(config["data"]["target_shape"]) or len(coarse) != len(target):
raise ValueError("Invalid CorrDiff NPZ tensor contract")
member_count = config["sampling"]["ensemble_size"]
ensemble = np.empty((member_count, len(coarse), *config["data"]["target_shape"]), dtype="float32")
options = {key: config["sampling"][key] for key in ("steps", "sigma_min", "sigma_max", "rho", "solver")}
batch_size = config["sampling"]["batch_size"]
with torch.inference_mode():
for member in range(member_count):
torch.manual_seed(config["seed"] + member)
for start in range(0, len(coarse), batch_size):
stop = min(start + batch_size, len(coarse))
batch = torch.from_numpy(coarse[start:stop]).to(device)
ensemble[member, start:stop] = model.sample(batch, **options).cpu().numpy()
del batch
output = ROOT / config["paths"]["predictions"]
output.parent.mkdir(parents=True, exist_ok=True)
np.savez_compressed(output, ensemble=ensemble, ensemble_mean=ensemble.mean(0),
ensemble_std=ensemble.std(0), target=target, protocol=np.asarray(protocol),
data_source=np.asarray(data_source), checkpoint=np.asarray(str(checkpoint_path)))
print(f"saved={output} ensemble={ensemble.shape} batch_size={batch_size}")
if __name__ == "__main__":
main()
|