使用 Python 和 C++ 并行 MCTS

简单介绍在 Python 和 C++ 中蒙特卡洛树搜索 (MCTS) 的并行化方法。

1 前言

当我们在 Python 中实现了 MCTS 后，可能会遇到性能问题，这时就需要将 MCTS 并行化。MCTS 的并行方法主要分为三种：

叶并行 (leaf parallelization), 即在叶节点扩展时进行并行。
根并行 (root parallelization), 即直接使用进程或线程创建多个不同的树，在不同的树中同时执行搜索。
树并行 (tree parallelization), 即多个线程在同一个树中进行并行，每个线程在树的不同部分执行搜索。

这些并行方法的详细介绍可以参考这篇论文的第 25 页。

图1 MCTS的并行化方法^[1]

2 叶并行

叶并行的方法比较简单，只需在遇到叶节点时，同时执行多次模拟 (simulation)，然后使用多次模拟的结果来代替之前的结果即可。但是在 AlphaZero 的算法中，叶并行的方法并不适用。因为该算法实际上并没有模拟 (simulation) 这个过程，而是使用神经网络直接返回当前节点的评估结果，因此需要使用其他并行方法。AlphaZero 的原理可以参考这篇文章。

3 根并行

根并行的实现同样比较简单，只要使用 Multiprocessing 或 concurrent.futures.ProcessPoolExecutor 直接调用 MCTS 即可。

import multiprocessing as mp

def mcts_main(a, b):
    '''
    mcts搜索的主函数
    输入参数：a, b
    '''
    ...

p = mp.Pool(10)           # 10个进程  
for i in range(50):       # 50个待执行任务
    p.apply_async(mcts_main, args=(a, b))
p.close()
p.join()

如果在 MCTS 中调用了 GPU 版的 Keras，需要设置显存按需增长，并且要将 import keras 放在函数内，否则将不能并行调用 model.predict()

import tensorflow as tf
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.Session(config=config)

def mcts_main(a, b):
    import keras            # import keras要放在函数内
    # mcts搜索的主函数
    ...

这样并行虽然简单，但缺点是程序占内存大，且显存容易不够用。

4 树并行

在 Python 中直接按照上述方法并行后，可能还是无法满足需求，需要更高效的并行方法。由于 GIL 的限制，在Python 中不能使用多线程来并行 MCTS，但实现树并行时需要进程/线程间较多的通信，如果直接在 Python 中用 multiprocessing 实现，可能带来较大的性能损失，同时实现起来也比较困难^[2]。

一种解决方法是用 C++ 实现树并行，再封装成 Python 接口，在 Python 中调用，这样就避开了 GIL 的限制。C++ 的树并行实现可以参考这个 Github 项目。这个项目使用的是 PyTorch，如果你使用的是 Keras，要在 C++ 中调用 Keras 可以参考这篇文章。将 C++ 封装成 Python 接口可以参考这篇文章。

封装成 Python 接口后，由于避开了 GIL 的限制，可以直接通过多线程调用。

import concurrent.futures

def mcts_main(a, b):
    '''
    mcts搜索的主函数
    输入参数：a, b
    '''
    ...

# 最大线程数10，50个待执行任务
with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
            futures = [executor.submit(mcts_main, a, b) for i in range(50)]
            for future in concurrent.futures.as_completed(futures):
                result = future.result()
                ...

5 参考资料

[1] Browne, C. B., Powley, E., Whitehouse, D., Lucas, S. M., Cowling, P. I., Rohlfshagen, P., … & Colton, S. (2012). A survey of monte carlo tree search methods. IEEE Transactions on Computational Intelligence and AI in games, 4(1), 1-43.

[2] https://stackoverflow.com/questions/52584142/mcts-tree-parallelization-in-python-possible

本人水平有限，如有不足之处，欢迎指出。