84 lines
2.9 KiB
Python
84 lines
2.9 KiB
Python
"""Profile GPU forward-pass throughput for the Deep CFR trainer network.
|
|
|
|
Builds the same DeepCFRMLP that ``DeepCFRTrainer.__init__`` constructs from
|
|
``legacy/deep-cfr/configs/default.yaml``, then measures average forward-pass time on
|
|
CUDA across a sweep of batch sizes. The goal is to decide whether batched
|
|
traversal inference (Optimization Priorities #5) is worth implementing.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from coolrl_lost_cities.games.classic.deep_cfr.encoding import input_dim
|
|
from coolrl_lost_cities.games.classic.game import GameState
|
|
|
|
from coolrl_lost_cities.games.classic.deep_cfr.config import load_config
|
|
from coolrl_lost_cities.games.classic.deep_cfr.networks import DeepCFRMLP
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parent.parent
|
|
CONFIG_PATH = REPO_ROOT / "legacy" / "deep-cfr" / "configs" / "default.yaml"
|
|
|
|
BATCH_SIZES = [1, 4, 16, 64, 256, 1024]
|
|
WARMUP_ITERS = 10
|
|
MEASURE_ITERS = 1000
|
|
|
|
|
|
def main() -> None:
|
|
if not torch.cuda.is_available():
|
|
raise SystemExit("CUDA is not available; this script requires a CUDA-capable GPU.")
|
|
|
|
cfg = load_config(CONFIG_PATH)
|
|
|
|
game_config = cfg.rules.to_lost_cities_config(seed=cfg.run.seed)
|
|
probe = GameState.new_game(game_config, seed=cfg.run.seed)
|
|
in_dim = input_dim(probe, cfg.encoding)
|
|
action_size = 2 * probe.config.hand_size + 1 + probe.config.n_colors
|
|
|
|
device = torch.device("cuda")
|
|
torch.manual_seed(cfg.run.seed)
|
|
network = DeepCFRMLP.from_config(in_dim, action_size, cfg.network).to(device)
|
|
network.eval()
|
|
|
|
print(
|
|
f"Network: DeepCFRMLP input_dim={in_dim} output_dim={action_size} "
|
|
f"hidden_size={cfg.network.hidden_size} num_layers={cfg.network.num_layers} "
|
|
f"activation={cfg.network.activation}"
|
|
)
|
|
print(f"Device: {torch.cuda.get_device_name(0)}")
|
|
print(f"Warmup iters: {WARMUP_ITERS} Measure iters: {MEASURE_ITERS}")
|
|
print()
|
|
|
|
results: list[tuple[int, float, float]] = []
|
|
with torch.inference_mode():
|
|
for bs in BATCH_SIZES:
|
|
x = torch.randn(bs, in_dim, device=device)
|
|
|
|
# Warm-up
|
|
for _ in range(WARMUP_ITERS):
|
|
network(x)
|
|
torch.cuda.synchronize()
|
|
|
|
start = time.perf_counter()
|
|
for _ in range(MEASURE_ITERS):
|
|
network(x)
|
|
torch.cuda.synchronize()
|
|
elapsed = time.perf_counter() - start
|
|
|
|
us_per_call = (elapsed / MEASURE_ITERS) * 1e6
|
|
us_per_state = us_per_call / bs
|
|
results.append((bs, us_per_call, us_per_state))
|
|
|
|
bs1_us_per_state = results[0][2]
|
|
print(f"{'batch_size':>10} | {'μs/call':>10} | {'μs/state':>10} | {'speedup_vs_bs1':>14}")
|
|
print("-" * 56)
|
|
for bs, us_call, us_state in results:
|
|
speedup = bs1_us_per_state / us_state
|
|
print(f"{bs:>10} | {us_call:>10.2f} | {us_state:>10.3f} | {speedup:>13.2f}x")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|