| |
| """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() |
|
|
