用PyTorch對Leela Zero進行神經網絡訓練

做者|Peter Yu
編譯|Flin
來源|towardsdatascience

最近，我一直在尋找方法來加快個人研究和管理個人實驗，特別是圍繞着寫訓練管道和管理實驗配置文件這兩個方面，我發現這兩個新項目叫作PyTorch Lightning和Hydra。PyTorch Lightning能夠幫助你快速編寫訓練管道，而Hydra能夠幫助你有效地管理配置文件。

• PyTorch Lightning：https://github.com/PyTorchLightning/pytorch-lightning
• Hydra：https://hydra.cc/

爲了練習使用它們，我決定爲Leela Zero(https://github.com/leela-zero/leela-zero) 編寫一個訓練管道。我這樣作，是由於這是一個範圍很廣的項目，涉及到使用多個gpu在大數據集上訓練大型網絡，能夠說是一個十分有趣的技術挑戰。此外，我之前曾經實現過一個更小版本的AlphaGo國際象棋（https://medium.com/@peterkeunwoo/beating-my-brother-in-chess-cb17739ffe2） ，因此我認爲這將是一個有趣的業餘項目。

在這個博客中，我將解釋這個項目的主要細節，以便你可以輕鬆理解我所作的工做。你能夠在這裏閱讀個人代碼:https://github.com/yukw777/leela-zero-pytorch

Leela Zero

第一步是找出Leela Zero神經網絡的內部工做原理。我大量引用了Leela Zero的文檔和它的Tensorflow訓練管道。

神經網絡結構

Leela Zero的神經網絡由一個殘差塔（ResNet 「tower」 ）組成，塔上有兩個「head」，即AlphaGo Zero論文（https://deepmind.com/blog/article/alphago-zero-starting-scratch） 中描述的負責策略的「頭」（policy head）和負責計算價值的「頭」（value head）。就像論文所述，策略「頭」和值「頭」開始的那幾個卷積濾波器都是1x1，其餘全部的卷積濾波器都是3x3。遊戲和棋盤特徵被編碼爲[批次大小，棋盤寬度，棋盤高度，特徵數量]形狀的張量，首先經過殘差塔輸入。而後，塔提取出抽象的特徵，並經過每一個「頭」輸入這些特徵，以計算下一步棋的策略機率分佈和遊戲的價值，從而預測遊戲的獲勝者。

你能夠在下面的代碼片斷中找到網絡的實現細節。

權重格式

Leela Zero使用一個簡單的文本文件來保存和加載網絡權重。文本文件中的每一行都有一系列數字，這些數字表示網絡的每一層的權重。首先是殘差塔，而後是策略頭，而後是值頭。

卷積層有2個權重行:

1. 與[output, input, filter size, filter size]形狀的卷積權值

2. 通道的誤差

Batchnorm層有2個權重行:

1. Batchnorm平均值

2. Batchnorm方差

內積(徹底鏈接)層有2個權重行:

1. 帶有[output, input]形狀的層權重

2. 輸出誤差

我編寫了單元測試來確保個人權重文件是正確的。我使用的另外一個簡單的完整性檢查是計算層的數量，在加載個人權值文件後，將其與Leela Zero進行比較。層數公式爲:

n_layers = 1 (version number) +
           2 (input convolution) + 
           2 (input batch norm) + 
           n_res (number of residual blocks) * 
           8 (first conv + first batch norm + 
              second conv + second batch norm) + 
           2 (policy head convolution) + 
           2 (policy head batch norm) + 
           2 (policy head linear) + 
           2 (value head convolution) + 
           2 (value head batch norm) + 
           2 (value head first linear) + 
           2 (value head second linear)

到目前爲止，這看起來很簡單，可是你須要注意一個實現細節。Leela Zero實際上使用卷積層的誤差來表示下一個歸一化層（batch norm）的可學習參數(gamma和beta)。這樣作是爲了使權值文件的格式(只有一行表示層權值，另外一行表示誤差)在添加歸一化層時沒必要更改。

目前，Leela Zero只使用歸一化層的beta項，將gamma設置爲1。那麼，實際上咱們該如何使用卷積誤差，來產生與在歸一化層中應用可學習參數相同的結果呢?咱們先來看看歸一化層的方程:

y = gamma * (x — mean)/sqrt(var — eps) + beta

因爲Leela Zero將gamma設爲1，則方程爲:

y = (x — mean)/sqrt(var — eps) + beta

如今，設定x_conv是沒有誤差的卷積層的輸出。而後，咱們想給x_conv添加一些誤差，這樣當你在沒有beta的歸一化層中運行它時，結果與在只有beta的歸一化層方程中運行x_conv是同樣的:

(x_conv + bias — mean)/sqrt(var — eps) = 
(x_conv — mean)/sqrt(var — eps) + beta 
x_conv + bias — mean = 
x_conv — mean + beta * sqrt(var — eps) 
bias = beta * sqrt(var — eps)

所以，若是咱們在權值文件中將卷積誤差設置爲beta * sqrt(var - eps)，咱們就會獲得指望的輸出，這就是LeelaZero所作的。

那麼，咱們如何實現它呢?在Tensorflow中，你能夠經過調用tf.layers.batch_normalization(scale=False)來告訴歸一化層要忽略gamma項，而後使用它。

遺憾的是，在PyTorch中，你不能將歸一化層設置爲只忽略gamma，你只能經過將仿射參數設置爲False: BatchNorm2d(out_channels, affine=False)，來忽略gamma和beta。因此，我把歸一化層設爲兩個都忽略，而後簡單地在後面加上一個張量，它表示beta。而後，使用公式bias = beta * sqrt(var - eps)來計算權值文件的卷積誤差。

訓練管道

在弄清了Leela Zeros的神經網絡的細節以後，就到了處理訓練管道的時候了。正如我提到的，我想練習使用兩個工具：PyTorch Lightning和Hydra，來加快編寫訓練管道和有效管理實驗配置。讓咱們來詳細瞭解一下我是如何使用它們的。

PyTorch Lightning

編寫訓練管道是我研究中最不喜歡的部分:它涉及大量重複的樣板代碼，並且很難調試。正由於如此，PyTorch Lightning對我來講就像一股清流，它是一個輕量級的庫，PyTorch沒有不少輔助抽象，在編寫訓練管道時，它負責處理大部分樣板代碼。它容許你關注你的訓練管道中更有趣的部分，好比模型架構，並使你的研究代碼更加模塊化和可調試。此外，它還支持多gpu和TPU的開箱即用訓練!

爲了使用PyTorch Lightning做爲個人訓練管道，我須要作的最多的編碼就是編寫一個類，我稱之爲NetworkLightningModule，它繼承自LightningModule來指定訓練管道的細節，並將其傳遞給訓練器。有關如何編寫本身的LightningModule的詳細信息，能夠參考PyTorch Lightning的官方文檔。

Hydra

我一直在研究的另外一部分是實驗管理。當你進行研究的時候，你不可避免地要運行大量不一樣的實驗來測試你的假設，因此，以一種可擴展的方式跟蹤它們是很是重要的。到目前爲止，我一直依賴於配置文件來管理個人實驗版本，可是使用平面配置文件很快就變得難以管理。使用模板是這個問題的一個解決方案。然而，我發現模板最終也會變得混亂，由於當你覆蓋多個層的值文件來呈現你的配置文件時，很難跟蹤哪一個值來自哪一個值文件。

另外一方面，Hydra是一個基於組件的配置管理系統。與使用單獨的模板和值文件來呈現最終配置不一樣，你能夠組合多個較小的配置文件來組成最終配置。它不如基於模板的配置管理系統靈活，但我發現基於組件的系統在靈活性和可維護性之間取得了很好的平衡。Hydra就是這樣一個專門爲研究腳本量身定作的系統。它的調用有點笨拙，由於它要求你將它用做腳本的主要入口點，但實際上我認爲有了這種設計，它很容易與你的訓練腳本集成。此外，它容許你經過命令行手動覆蓋配置，這在運行你的實驗的不一樣版本時很是有用。我經常使用Hydra管理不一樣規模的網絡架構和訓練管道配置。

評估

爲了評估個人訓練網絡，我使用GoMill（https://github.com/mattheww/gomill） 來舉行圍棋比賽。它是一個運行在Go Text Protocol (GTP)引擎上的比賽的庫，Leela Zero就是其中之一。你能夠在這裏（https://github.com/yukw777/leela-zero-pytorch/blob/master/eval/bg-vs-sm.ctl） 找到我使用的比賽配置。

結論

經過使用PyTorch-Lightning和Hydra，可以極大地加快編寫訓練管道的速度，並有效地管理實驗配置文件。我但願這個項目和博客文章也能對你的研究有所幫助。你能夠在這裏查看代碼:https://github.com/yukw777/leela-zero-pytorch

原文連接：https://towardsdatascience.com/training-neural-networks-for-leela-zero-using-pytorch-and-pytorch-lightning-bbf588683065

歡迎關注磐創AI博客站：
http://panchuang.net/

sklearn機器學習中文官方文檔：
http://sklearn123.com/

歡迎關注磐創博客資源彙總站：
http://docs.panchuang.net/