前言
專案地址：github.com/hijkzzz/alp… 如果覺得本專案還可以，請贊顆星星支援一下...感謝
AlphaZero演算法已經發布了一年多了，GitHub也有各種各樣的實現，有一千行Python程式碼單執行緒低效能版，也有數萬行C++程式碼的分散式版本。但是這些實現都不能滿足一般的演算法愛好者的需求，即一個簡單的並且單機的可執行的高效能AlphaZero演算法。
一圖解密AlphaZero
首先我們先了解一下AlphaZero演算法的原理
可以看到AlaphaGo Zero的演算法流程分為：自對弈（利用蒙特卡洛樹搜尋）N局生成棋譜 ==> 利用生成的棋譜訓練網路 ==> 評估新訓練的網路
分析
對於Python版本的AlphaZero演算法，通常受限制於GIL，過程中最耗時間的自對弈階段（見下圖）無法並行化，所以最直接的優化方式是使用C++這種高效能語言實現底層運算細節，用Python封裝。
解決方法
執行緒池
為了並行化自對弈過程，首先我們需要實現一個C++的執行緒池。關於執行緒池網上有很多的資料可以參考，這裡就不多做敘述。
Root Parallelization
從演算法流程圖中可以看到，自對弈過程使用蒙特卡洛樹搜尋實現，所以有兩個維度可以並行化自對弈：Root Parallelization和Tree Parallelization。其中Root Parallelization指的是同時開啟N局對弈，每個執行緒負責一局遊戲。Tree Parallelization指的是把單局遊戲中的蒙特卡洛樹搜尋（MCTS）並行化。於是用N個執行緒就很容易實現Root Parallelization，下面我們討論Tree Parallelization。
Tree Parallelization
首先分析一下蒙特卡洛樹搜尋（MCTS）的執行過程：
每執行一步棋子，MCTS要執行M次落子模擬，每次模擬就是一次遞迴過程，如下：
-
Select，如果當前節點不是葉子節點則通過特定的UCT演算法（探索-利用演算法，通過神經網路預測的勝率值（q值）以及先驗概率計算選擇概率，勝率/先驗概率越高選擇機率越大）找出最優的下一個落子位置，搜尋進入下一層，直到當前節點是葉子節點。
-
Expand and evaluate，如果當前節點是葉子節點，這裡分為兩種情況：
- 當前節點遊戲結束，某一方獲勝，則進行Backup向上回溯更新父節點的勝率值
- 如果遊戲沒有結束，則用神經網路預測當前節點的勝率和下一層的先驗概率，用這個先驗概率展開此節點，然後進行Backup向上回溯更新父節點的勝率值（q值）
-
Backup，每個節點儲存一個勝率值（q值），q值等於贏的次數/訪問次數，backup從結束狀態向上更新這個值以及訪問次數。
-
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
最後我們把上述相關的C++程式碼用SWIG封裝成Python介面，以供主程式呼叫。雖然這會導致一部分效能開銷，但是大大提高了開發的效率。
效果
經過測試，並行化後的訓練效率至少提升了10倍。簡單的計算一下，假設每個MCTS4個執行緒，同時玩4局遊戲，即4x4=16倍，考慮鎖和緩衝佇列以及Python介面的開銷，提升數量級是合理的。此外只要GPU足夠強悍，提升執行緒數還能繼續提高效能。最後我用了12個小時在一塊GTX1070上訓練了一個標準的15x15的五子棋演算法，已足夠可以對我的棋藝進行碾壓。
參考文獻
- Mastering the Game of Go without Human Knowledge
- Mastering Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithm
- Parallel Monte-Carlo Tree Search
- An Analysis of Virtual Loss in Parallel MCTS
- A Lock-free Multithreaded Monte-Carlo Tree Search Algorithm