我正在實現一個 AlphaZero 版本（AlphaGo 的最新版本），以應用于其他領域。該算法的關鍵是狀態空間 (CPU) 的蒙特卡洛樹搜索與來自 eval 模式 (GPU) 的神經網絡的“直覺”（概率）交錯。然后使用 MCTS 結果來訓練神經網絡。我已經通過啟動多個進程來并行化 CPU 執行，每個進程都建立自己的樹。這是有效的，現在已經導致GPU 瓶頸！（nvidia-smi 始終以 100% 顯示 GPU）我設計了 2 種策略來并行化 GPU 評估，但是它們都有問題。每個進程僅在其自己的樹中的批次上評估網絡。在我最初的幼稚實現中，這意味著批量大小為 1。但是，通過重構一些代碼并添加“虛擬損失”來阻止（但不完全阻止）同一節點被選中兩次，我們可以獲得更大的批量大小 1 -4。這里的問題是，在我們評估批次或準確性受到影響之前，我們不能允許大的延遲，所以小批量是這里的關鍵。將批次發送到一個中央“神經網絡工作者”線程，該線程組合和評估它們。這可以在 32 個或更多的大批量中完成，因此可以非常有效地使用 GPU。這里的問題是樹工作者發送 CUDA 張量“往返”，這不受 PyTorch 支持。如果我先克隆它們是受支持的，但是所有不斷的復制使這種方法比第一種方法慢。我在想也許一個我沒有看到的聰明的批處理方案可以使第一種方法起作用。使用多個 GPU 也可以加速第一種方法，但是 PyTorch 本身并不支持我想要的那種并行性。也許將所有張量保留在 NN worker 中并且只發送 id 可以改進第二種方法，但是這里的困難是如何有效地同步以獲得大批量而不使 CPU 線程等待太久。我幾乎沒有發現有關 AlphaZero 或 AlphaGo Zero 在各自論文中如何并行化的信息。我能夠在網上找到有限的信息，但這導致我改進了第一種方法。如果有任何建議，我將不勝感激，特別是如果我錯過了某些觀點或方法。