我正在实现一个AlphaZero（AlphaGo的最新版本）的版本，应用于其他领域。

该算法的核心是蒙特卡洛树搜索（CPU）与神经网络在评估模式下的“直觉”（概率，GPU）交替进行。随后使用MCTS结果来训练神经网络。

我已经通过启动多个进程并行化了CPU执行，每个进程构建自己的树。这很有效，现在导致了GPU瓶颈！（nvidia-smi显示GPU始终处于100%）

我已经设计了两种策略来并行化GPU评估，但它们都有问题。

• 每个进程仅评估来自其自身树的批次。 在我最初的简单实现中，这意味着批次大小为1。然而，通过重构一些代码并添加“虚拟损失”来阻止（但不完全阻塞）同一个节点被选中两次，我们可以得到1-4的大小批次。这里的问题是，我们不能允许在评估批次之前有大的延迟，否则准确性会受到影响，因此小批次大小在这里是关键。

• 将批次发送到中央“神经网络工作者”线程，该线程合并并评估它们。 这可以以32或更大的大批次进行，因此可以非常有效地使用GPU。这里的问题是树工作者发送CUDA张量进行“往返”，这不被PyTorch支持。如果我先克隆它们，这是支持的，但所有这些持续的复制使得这种方法比第一种方法慢。

我在想，也许一个我没有想到的巧妙的批处理方案可以使第一种方法奏效。使用多个GPU也可以加速第一种方法，但PyTorch原生不支持我想要的那种并行性。也许将所有张量保存在NN工作者中，仅发送ID可能会改善第二种方法，但这里的难点是如何有效地同步以获得大批次，而不让CPU线程等待太长时间。

我在各自的论文中几乎找不到关于AlphaZero或AlphaGo Zero如何并行化的信息。然而，我在网上找到了有限的信息，这让我改进了第一种方法。

如果有任何建议，特别是如果有我遗漏的要点或方法，我将不胜感激。

回答：

以tensorflow serving为例，预测服务可以运行在不同的进程中，运行一个服务来接收来自工作者（运行MCTS进程并向预测服务发送预测请求）的请求。我们可以保持一个从套接字地址到套接字本身的字典。

预测服务可以读取每个查询的正文及其头部（每个查询的头部不同），我们可以将这些头部放入队列中。在等待最多100毫秒或当前批次大于批次大小的情况下，预测运行。在GPU给出结果后，我们遍历结果，因为顺序与队列中的头部相同，我们可以通过套接字根据每个头部（可以从我们上面保持的字典中查找）发送响应。

由于每个查询都带有不同的头部，你不能错过请求、响应和套接字。虽然你可以运行一个带有GPU卡的tensorflow serving，同时运行多个工作者，以保持批次大小足够大，从而获得更大的吞吐量。

我还在这里找到了一个批处理机制：https://github.com/richemslie/galvanise_zero