| """Train a pure NumPy RF; torchrun ranks build disjoint tree subsets.""" |
|
|
| import json |
| import os |
| import pickle |
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
| import yaml |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
| from model.meteonorm_rf import (FEATURE_NAMES, FORMAT_VERSION, MODEL_NAME, |
| MultiOutputRandomForest, MeteoNormRF, |
| encode_features, merge_states, save_checkpoint) |
|
|
|
|
| def main(): |
| config = yaml.safe_load((ROOT / "conf/config.yaml").read_text()) |
| seed = int(config["seed"]) |
| data = np.load(ROOT / config["data"]["path"]) |
| x, y = encode_features(data), data["pollution"].astype(np.float32) |
| rng = np.random.default_rng(seed) |
| order = rng.permutation(len(x)) |
| cut = int(float(config["data"]["train_fraction"]) * len(order)) |
| train_indices, test_indices = order[:cut], order[cut:] |
| world, rank = int(os.environ.get("WORLD_SIZE", "1")), int(os.environ.get("RANK", "0")) |
| options = config["model"]["engineering"] |
| trees = int(options["trees"]) |
| assigned = list(range(rank, trees, world)) |
| forest = MultiOutputRandomForest(n_trees=trees, seed=seed, |
| max_depth=int(options["max_depth"]), |
| min_samples_leaf=int(options["min_samples_leaf"]), |
| max_features=options["max_features"], |
| split_candidates=int(options["split_candidates"])) |
| model = MeteoNormRF(forest).fit(x[train_indices], y[train_indices], assigned) |
| checkpoint = ROOT / config["paths"]["checkpoint"] |
| checkpoint.parent.mkdir(parents=True, exist_ok=True) |
| shard = checkpoint.with_suffix(f".rank{rank}.pkl") |
| with open(shard, "wb") as stream: |
| pickle.dump(model.state_dict(), stream) |
| if world > 1: |
| import torch |
| torch.distributed.init_process_group("gloo") |
| torch.distributed.barrier() |
| if rank == 0: |
| states = [] |
| for item in range(world): |
| with open(checkpoint.with_suffix(f".rank{item}.pkl"), "rb") as stream: |
| states.append(pickle.load(stream)) |
| merged = MeteoNormRF.from_state_dict(merge_states(states)) |
| metadata = {"feature_names": FEATURE_NAMES, "train_indices": train_indices, |
| "test_indices": test_indices, "split": "seeded random 70/30", |
| "paper_trees": config["paper_protocol"]["trees"], "engineering_trees": trees} |
| save_checkpoint(checkpoint, merged, metadata) |
| metrics = ROOT / config["paths"]["training_metrics"] |
| metrics.parent.mkdir(parents=True, exist_ok=True) |
| metrics.write_text(json.dumps({"format_version": FORMAT_VERSION, "model": MODEL_NAME, |
| "rows": len(x), "train_rows": cut, "test_rows": len(x) - cut, |
| "trees": trees, "world_size": world}, indent=2) + "\n") |
| for item in range(world): |
| checkpoint.with_suffix(f".rank{item}.pkl").unlink() |
| print(f"checkpoint={checkpoint.relative_to(ROOT)} trees={trees} train={cut} test={len(x)-cut}") |
| if world > 1: |
| torch.distributed.barrier() |
| torch.distributed.destroy_process_group() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
