File size: 3,330 Bytes
b11ef36 | 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 | """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()
|