File size: 2,841 Bytes
73ddb67 | 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 | #!/usr/bin/env python3
"""Run all serialized trees and retain ensemble and per-tree predictions."""
from __future__ import annotations
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 / "model"))
from ml_modis import BootstrapRandomForestRegressor, validate_multimodal_keys
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--config", default=str(ROOT / "conf/config.yaml"))
parser.add_argument("--data", default=None)
parser.add_argument("--checkpoint", default=None)
parser.add_argument("--output", default=None)
args = parser.parse_args()
config = yaml.safe_load(Path(args.config).read_text())
with np.load(ROOT / (args.data or config["data"]["path"])) as archive:
data = {key: archive[key] for key in archive.files}
validate_multimodal_keys(data)
checkpoint = torch.load(ROOT / (args.checkpoint or config["paths"]["checkpoint"]), map_location="cpu", weights_only=False)
if checkpoint.get("format_version") != config["format_version"]:
raise ValueError("Checkpoint format_version does not match configuration")
targets = list(checkpoint["model_config"]["targets"])
tree_count = len(next(iter(checkpoint["model"].values()))["state"]["trees"])
tree_predictions = np.full((data["X"].shape[0], len(targets), tree_count), np.nan, dtype=np.float32)
for month in checkpoint["model_config"]["months"]:
mask = data["month"] == month
for target_index, target in enumerate(targets):
model = BootstrapRandomForestRegressor.from_state_dict(checkpoint["model"][f"{month}:{target}"]["state"])
tree_predictions[mask, target_index, :] = model.predict_trees(data["X"][mask])
prediction = tree_predictions.mean(axis=2)
safe_prediction = np.where(np.abs(prediction) > 1e-8, prediction, np.nan)
ratio = data["Y"] / safe_prediction
if not np.isfinite(prediction).all() or not np.isfinite(ratio).all():
raise FloatingPointError("Inference produced non-finite values")
output = ROOT / (args.output or config["paths"]["predictions"])
output.parent.mkdir(parents=True, exist_ok=True)
np.savez_compressed(output, pred=prediction, pred_trees=tree_predictions, obs=data["Y"],
obs_over_pred=ratio, relative_response=ratio - 1.0,
year=data["year"], month=data["month"], platform=data["platform"],
latitude=data["latitude"], longitude=data["longitude"],
target_names=np.asarray(targets))
print(f"output={output.relative_to(ROOT)} samples={prediction.shape[0]} "
f"targets={targets} trees_per_prediction={tree_count}")
if __name__ == "__main__":
main()
|