File size: 4,612 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 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 | #!/usr/bin/env python3
"""Train month-by-target ML-MODIS forests, optionally task-parallel under torchrun."""
from __future__ import annotations
import argparse
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 / "model"))
from ml_modis import BootstrapRandomForestRegressor, feature_names, regression_metrics, validate_multimodal_keys
def args_parser() -> argparse.Namespace:
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("--paper-model", action="store_true")
parser.add_argument("--trees", type=int, default=None)
return parser.parse_args()
def distributed_context() -> tuple[int, int]:
world = int(os.environ.get("WORLD_SIZE", "1"))
rank = int(os.environ.get("RANK", "0"))
if world > 1:
torch.distributed.init_process_group(backend="gloo")
return rank, world
def main() -> None:
args = args_parser()
config = yaml.safe_load(Path(args.config).read_text())
settings = dict(config["model"])
if args.paper_model:
settings.update(config["paper_model"])
if args.trees is not None:
settings["trees"] = args.trees
data_path = ROOT / (args.data or config["data"]["path"])
with np.load(data_path) as archive:
data = {key: archive[key] for key in archive.files}
validate_multimodal_keys(data)
rank, world = distributed_context()
months = config["data"]["months"]
targets = config["data"]["variables"]["targets"]["names"]
tasks = [(int(month), target_index, target) for month in months
for target_index, target in enumerate(targets)]
local_models = {}
for task_index, (month, target_index, target) in enumerate(tasks):
if task_index % world != rank:
continue
mask = (data["month"] == month) & (data["year"] != config["train"]["excluded_year"])
x, y = data["X"][mask], data["Y"][mask, target_index]
model = BootstrapRandomForestRegressor(
n_trees=int(settings["trees"]), min_leaf=int(settings["min_leaf"]),
max_features=int(settings["max_features"]), bootstrap_fraction=float(settings["bootstrap_fraction"]),
max_depth=settings["max_depth"], split_candidates=int(settings["split_candidates"]),
seed=int(config["runtime"]["seed"] + task_index * 1009),
).fit(x, y)
oob, counts = model.oob_predict(x)
local_models[f"{month}:{target}"] = {
"state": model.state_dict(), "oob_metrics": regression_metrics(y[counts > 0], oob[counts > 0]),
"train_samples": int(mask.sum()), "excluded_year": int(config["train"]["excluded_year"]),
}
print(f"rank={rank} trained month={month} target={target} samples={mask.sum()}", flush=True)
if world > 1:
gathered = [None] * world if rank == 0 else None
torch.distributed.gather_object(local_models, gathered, dst=0)
if rank == 0:
local_models = {key: value for shard in gathered for key, value in shard.items()}
if rank == 0:
checkpoint = ROOT / (args.checkpoint or config["paths"]["checkpoint"])
checkpoint.parent.mkdir(parents=True, exist_ok=True)
model_config = {
"architecture": "BootstrapRandomForestRegressor", "settings": settings,
"targets": targets, "months": months, "input_features": 114,
"feature_names": feature_names(), "excluded_year": int(config["train"]["excluded_year"]),
}
torch.save({"model": local_models, "model_config": model_config,
"format_version": config["format_version"],
"training": {"paper_model": args.paper_model, "distributed_world_size": world}}, checkpoint)
summary = {key: value["oob_metrics"] for key, value in sorted(local_models.items())}
metrics_path = ROOT / config["paths"]["training_metrics"]
metrics_path.parent.mkdir(parents=True, exist_ok=True)
metrics_path.write_text(json.dumps({"format_version": config["format_version"],
"models": summary}, indent=2, allow_nan=False) + "\n")
print(json.dumps({"checkpoint": str(checkpoint), "models": len(local_models), "oob": summary}, indent=2))
if world > 1:
torch.distributed.destroy_process_group()
if __name__ == "__main__":
main()
|