五子棋 AlphaZero（PyTorch GPU 版）

一个可自我对弈训练的迷你 AlphaZero，包含五子棋规则、MCTS、PyTorch 残差卷积策略价值网络、棋盘对称增强、经验回放、checkpoint 和命令行人机对战。

网络输入有三个平面：当前方棋子、对手棋子、最后一手位置。第三平面不再直接暴露黑白颜色，避免Value仅凭颜色拟合某一方固定获胜的训练捷径。

训练时会并行推进一轮中的所有自我对弈，把多个MCTS叶子合成一次GPU批量推理，并在真实落子后复用已搜索的子树。--games 20 因此也代表推理批次最多为20；通常比逐局、单样本推理快得多。

MCTS在每个展开节点先执行确定性战术检查：如果当前方能一步成五，只保留获胜点；否则如果对手下一步能成五，只保留必须封堵的位置。这能阻止弱网络和低搜索预算共同漏掉最基本的一步杀。

CPU热点采用Numba编译，MCTS节点使用225维连续数组并只在动作真正被访问时创建子节点。另提供可选的CPU规则进程池；由于Numba规则扫描已经很快，小批量下进程通信可能反而更慢，因此默认关闭。

环境

当前机器使用：

• Python 3.8
• PyTorch 2.4.1 + CUDA 12.4
• NVIDIA RTX 2000 Ada 16GB

进入项目并激活环境：

cd "<项目所在目录>\gomoku-alpha-zero"
conda activate py38

验证 CUDA：

python -c "import torch; print(torch.__version__, torch.cuda.is_available(), torch.cuda.get_device_name(0))"

快速验证

python -m unittest discover -s tests -v
python main.py train --size 5 --model models/test_5x5.pt --channels 16 --blocks 1 --iterations 1 --games 1 --simulations 5 --batches 2

训练 9×9

先用较小配置确认流程：

python main.py train --iterations 2 --games 3 --simulations 30 --batches 10

正式训练：

python main.py train --iterations 100 --games 20 --simulations 100 --batches 50 --channels 64 --blocks 4

默认保存到 models/gomoku_9x9.pt。每轮结束保存一次；再次运行相同命令会自动载入模型和 Adam 优化器状态继续训练。

训练标准 15×15

python main.py train `
  --size 15 `
  --model models/gomoku_15x15_lastmove.pt `
  --channels 64 `
  --blocks 5 `
  --iterations 100 `
  --games 20 `
  --simulations 150 `
  --batches 60

RTX 2000 Ada 16GB 可以使用更大的网络，例如 --channels 96 --blocks 6，但自我对弈中的MCTS是串行单局面推理，GPU显存并不是主要瓶颈。

人机对战

python main.py play --model models/gomoku_9x9.pt --simulations 300
python main.py play --model models/gomoku_15x15_lastmove.pt --simulations 500 --color white

棋盘中 X 是黑棋，O 是白棋。输入 E5，或输入 5 5（行、列从1开始）。输入 q 退出。

仓库包含当前最新的15×15训练模型 models/gomoku_15x15_lastmove.pt，可以直接使用上面的命令对弈，也可以在此基础上继续训练。

参数

• --channels：卷积通道数，越大容量越强、计算越慢，仅创建新模型时生效。
• --blocks：残差块数，仅创建新模型时生效。
• --simulations：每一步MCTS模拟次数。
• --games：每轮自我对弈局数。
• --iterations：训练轮数。
• --batches：每轮网络更新次数。
• --batch-size：每次网络更新的样本数。
• --device auto|cuda|cpu：默认自动使用CUDA。
• --workers：CPU规则进程数，默认0关闭。只有在较大的games配置下实测吞吐更高时才建议设置，例如--workers 4。
• --temperature-moves：每局会在0到该值之间随机选择前期采样手数，默认上限为12。例如不同棋局可能分别采样前0、3、7或12手；超过各自截止手数后选择MCTS访问次数最多的位置。
• --prior-power：MCTS使用Policy前进行幂变换。训练默认0.5（开根号并重新归一化），会压低热门位置、提高低概率位置的探索机会；1.0表示不变。一步胜和一步防候选始终优先，不受该参数裁剪。
• 默认对每轮生成的完整落子序列去重：相同棋谱只向经验池加入一次，但胜负率和平均手数仍统计全部棋局。日志会显示唯一棋局和去重数量。调试时可用--keep-duplicate-games关闭去重。

比较Worker效果时应使用相同配置测量“总手数÷耗时”，不要只看任务管理器CPU占用。本机8局短基准中，关闭Worker比4个Worker更快，因为Numba规则计算耗时已经低于Windows进程通信开销。

旧版 .npz 是单隐藏层NumPy模型，不能继续使用；之前使用“当前是否黑棋”作为第三平面的 .pt 模型也会被输入编码校验拒绝。请使用新的模型文件名重新训练，旧文件会保留且不会被覆盖。

当前限制

• 经验回放池只保存在当前进程内，重启不会恢复。
• 一轮内的自我对弈会批量搜索；不同迭代之间仍然是串行的。
• 规则为自由五子棋，没有三三、四四和长连禁手。
• 人机对战入口仍使用简单的单局搜索；训练路径会复用子树。