File size: 1,835 Bytes
b97f43c 07f4740 b97f43c 07f4740 b97f43c 07f4740 | 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 | #!/usr/bin/env python3
"""Find the best (workers x threads) for self-play on this machine."""
import argparse, json, os, sys, time
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import torch
from tinychess.config import REGISTRY
from tinychess.model import TinyChess
from tinychess.selfplay import SelfPlayConfig
from tinychess.parallel import play_games_parallel
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--games", type=int, default=128)
ap.add_argument("--max-plies", type=int, default=80)
ap.add_argument("--grid", default="1x1,8x2,16x2,32x2,48x2,64x2,96x2,48x4")
ap.add_argument("--out", default="runs/bench_parallel.json")
a = ap.parse_args()
m = TinyChess(REGISTRY["baseline"]()); m.eval()
res = {"cpu_count": os.cpu_count(), "games": a.games, "runs": []}
for spec in a.grid.split(","):
w, t = (int(x) for x in spec.split("x"))
cfg = SelfPlayConfig(n_games=a.games, max_plies=a.max_plies, steps=4,
adjudicate="material", seed=0)
t0 = time.time()
recs, st = play_games_parallel(m, None, cfg, n_workers=w, threads_per_worker=t)
dt = time.time() - t0
moves = sum(len(r.moves) for r in recs)
row = {"workers": w, "threads": t, "sec": round(dt, 2),
"moves_per_s": round(moves / dt, 1), "games": len(recs)}
res["runs"].append(row)
print(f"{w:3d}w x {t}t {dt:7.2f}s {row['moves_per_s']:8.1f} moves/s", flush=True)
json.dump(res, open(a.out, "w"), indent=2)
best = max(res["runs"], key=lambda r: r["moves_per_s"])
res["best"] = best
json.dump(res, open(a.out, "w"), indent=2)
print(f"\nBEST: {best['workers']}w x {best['threads']}t = {best['moves_per_s']} moves/s")
if __name__ == "__main__":
main()
|