动手实现并行版AlphaZero五子棋

前言

项目连接：github.com/hijkzzz/alp…

AlphaZero算法已经发布了一年多了，GitHub也有各类各样的实现，有一千行Python代码单线程低性能版，也有数万行C++代码的分布式版本。可是这些实现都不能知足通常的算法爱好者的需求，即一个简单的而且单机的可运行的高性能AlphaZero算法。

一图解密AlphaZero

首先咱们经过一张图了解一下AlphaZero算法的原理

能够看到AlaphaGo Zero的算法流程分为：

1. 自对弈（利用蒙特卡洛树搜索）N局生成棋谱
2. 利用生成的棋谱训练网络
3. 评估新训练的网络

分析

对于Python版本的AlphaZero算法，一般受限制于GIL，过程当中最耗时间的自对弈阶段（见下图）没法并行化，因此最直接的优化方式是使用C++这种高性能语言实现底层运算细节。

解决方法

线程池

github.com/hijkzzz/alp…

为了并行化自对弈过程，首先咱们须要实现一个C++的线程池。关于线程池网上有不少的资料能够参考，这里就很少作叙述。

Root Parallelization

从算法流程图中能够看到，自对弈过程使用蒙特卡洛树搜索实现，因此有两个维度能够并行化自对弈：Root Parallelization和Tree Parallelization。其中Root Parallelization指的是同时开启N局对弈，每一个线程负责一局游戏。Tree Parallelization指的是把单局游戏中的蒙特卡洛树搜索（MCTS）并行化。因而用N个线程就很容易实现Root Parallelization，下面咱们讨论Tree Parallelization。

Tree Parallelization

首先分析一下蒙特卡洛树搜索（MCTS）的运行过程：

每执行一步棋子，MCTS要执行M次落子模拟，每次模拟就是一次递归过程，以下：

1. Select，若是当前节点不是叶子节点则经过特定的UCT算法（探索-利用算法，经过神经网络预测的胜率值（q值）以及先验几率计算选择几率，胜率/先验几率越高选择概率越大）找出最优的下一个落子位置，搜索进入下一层，直到当前节点是叶子节点。

2. Expand and evaluate，若是当前节点是叶子节点，这里分为两种状况：

   • 当前节点游戏结束，某一方获胜，则进行Backup向上回溯更新父节点的胜率值
   • 若是游戏没有结束，则用神经网络预测当前节点的胜率和下一层的先验几率，用这个先验几率展开此节点，而后进行Backup向上回溯更新父节点的胜率值（q值）

3. Backup，每一个节点保存一个胜率值（q值），q值等于赢的次数/访问次数，backup从结束状态向上更新这个值以及访问次数。

4. Play，实际游戏中落子的时候选择根节点下访问次数最多的子节点便可（由于q值越大的节点select的几率越大，访问次数也越多）。

因此咱们能够同时进行M'（小于M）次模拟，因此对一些关键数据就要加锁，好比蒙特卡洛树的父子节点关系，访问次数，q值等。也有人研发出了一些无锁的算法[5]，可是由于预先分配树节点的关系，对内存的占用量极大，通常的机器跑不起来，因此这里用的是加锁版的并行蒙特卡洛树搜索。

Virtual Loss

对于Tree Parallelization，若是咱们简单的把蒙特卡洛搜索（MCTS）并行化，那么会遇到一个问题：M'个线程常常会搜索同一个节点，这样咱们的并行化就失去了意义，由于搜索同一个节点意味着重复工做。因此在UCT算法中，当一个节点被一个线程访问时，咱们加入一个Virtual Loss的惩罚，这样其它线程就不太可能会选择这个节点进行搜索。

LibTorch

由于MCTS的过程当中须要用到神经网络预测胜率和先验几率，因此C++须要调用Python实现的神经网络预测方法，可是这样又会回到原点。即Pyhton的GIL限制会致使并行化的自对弈被强制串行化执行。因此咱们使用PyTorch的C++前端LibTorch实现神经网络。

CUDA Stream

工做后对于运行在GPU上的神经网络来讲，实际上咱们的程序仍是没有真正的并行化。这是由于LibTorch的预测执行受限制于Default CUDA Steam，即默认是串行的，会致使多线程调用预测被阻塞。有两个方法来避免这个问题：1. 用多个CUDA Stream 2.合并预测请求。这里咱们使用的方法是用缓冲队列合并多个预测，一次性推送到GPU，这样就防止了GPU工做流的争用致使线程阻塞。

SWIG

www.swig.org/tutorial.ht…

最后咱们把上述相关的C++代码用SWIG封装成Python接口，以供主程序调用。虽然这会致使一部分性能开销，可是大大提升了开发的效率。

测试

通过测试，并行化后的训练效率至少提高了10倍。简单的计算一下，假设每一个MCTS用4个线程，同时玩4局游戏，即4x4=16倍，考虑锁和缓冲队列以及Python接口的开销，提高数量级是合理的。此外只要GPU足够强悍，提高线程数还能继续提升性能。最后我用了一天时间在一块GTX1070上训练了一个标准的15x15的五子棋算法，已经能够完败普通玩家。

参考文献

1. Mastering the Game of Go without Human Knowledge
2. Mastering Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithm
3. Parallel Monte-Carlo Tree Search
4. An Analysis of Virtual Loss in Parallel MCTS
5. A Lock-free Multithreaded Monte-Carlo Tree Search Algorithm