OneScience-Group/ML-MODIS

1#!/usr/bin/env python32"""Train month-by-target ML-MODIS forests, optionally task-parallel under torchrun."""3 4from __future__ import annotations5 6import argparse7import json8import os9import sys10from pathlib import Path11 12import numpy as np13import torch14import yaml15 16ROOT = Path(__file__).resolve().parents[1]17sys.path.insert(0, str(ROOT / "model"))18from ml_modis import BootstrapRandomForestRegressor, feature_names, regression_metrics, validate_multimodal_keys19 20 21def args_parser() -> argparse.Namespace:22    parser = argparse.ArgumentParser()23    parser.add_argument("--config", default=str(ROOT / "conf/config.yaml"))24    parser.add_argument("--data", default=None)25    parser.add_argument("--checkpoint", default=None)26    parser.add_argument("--paper-model", action="store_true")27    parser.add_argument("--trees", type=int, default=None)28    return parser.parse_args()29 30 31def distributed_context() -> tuple[int, int]:32    world = int(os.environ.get("WORLD_SIZE", "1"))33    rank = int(os.environ.get("RANK", "0"))34    if world > 1:35        torch.distributed.init_process_group(backend="gloo")36    return rank, world37 38 39def main() -> None:40    args = args_parser()41    config = yaml.safe_load(Path(args.config).read_text())42    settings = dict(config["model"])43    if args.paper_model:44        settings.update(config["paper_model"])45    if args.trees is not None:46        settings["trees"] = args.trees47    data_path = ROOT / (args.data or config["data"]["path"])48    with np.load(data_path) as archive:49        data = {key: archive[key] for key in archive.files}50    validate_multimodal_keys(data)51    rank, world = distributed_context()52    months = config["data"]["months"]53    targets = config["data"]["variables"]["targets"]["names"]54    tasks = [(int(month), target_index, target) for month in months55              for target_index, target in enumerate(targets)]56    local_models = {}57    for task_index, (month, target_index, target) in enumerate(tasks):58        if task_index % world != rank:59            continue60        mask = (data["month"] == month) & (data["year"] != config["train"]["excluded_year"])61        x, y = data["X"][mask], data["Y"][mask, target_index]62        model = BootstrapRandomForestRegressor(63            n_trees=int(settings["trees"]), min_leaf=int(settings["min_leaf"]),64            max_features=int(settings["max_features"]), bootstrap_fraction=float(settings["bootstrap_fraction"]),65            max_depth=settings["max_depth"], split_candidates=int(settings["split_candidates"]),66            seed=int(config["runtime"]["seed"] + task_index * 1009),67        ).fit(x, y)68        oob, counts = model.oob_predict(x)69        local_models[f"{month}:{target}"] = {70            "state": model.state_dict(), "oob_metrics": regression_metrics(y[counts > 0], oob[counts > 0]),71            "train_samples": int(mask.sum()), "excluded_year": int(config["train"]["excluded_year"]),72        }73        print(f"rank={rank} trained month={month} target={target} samples={mask.sum()}", flush=True)74    if world > 1:75        gathered = [None] * world if rank == 0 else None76        torch.distributed.gather_object(local_models, gathered, dst=0)77        if rank == 0:78            local_models = {key: value for shard in gathered for key, value in shard.items()}79    if rank == 0:80        checkpoint = ROOT / (args.checkpoint or config["paths"]["checkpoint"])81        checkpoint.parent.mkdir(parents=True, exist_ok=True)82        model_config = {83            "architecture": "BootstrapRandomForestRegressor", "settings": settings,84            "targets": targets, "months": months, "input_features": 114,85            "feature_names": feature_names(), "excluded_year": int(config["train"]["excluded_year"]),86        }87        torch.save({"model": local_models, "model_config": model_config,88                    "format_version": config["format_version"],89                    "training": {"paper_model": args.paper_model, "distributed_world_size": world}}, checkpoint)90        summary = {key: value["oob_metrics"] for key, value in sorted(local_models.items())}91        metrics_path = ROOT / config["paths"]["training_metrics"]92        metrics_path.parent.mkdir(parents=True, exist_ok=True)93        metrics_path.write_text(json.dumps({"format_version": config["format_version"],94                                            "models": summary}, indent=2, allow_nan=False) + "\n")95        print(json.dumps({"checkpoint": str(checkpoint), "models": len(local_models), "oob": summary}, indent=2))96    if world > 1:97        torch.distributed.destroy_process_group()98 99 100if __name__ == "__main__":101    main()102