1. 程式人生 > 實用技巧 >用PyTorch對Leela Zero進行神經網路訓練

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

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

最近,我一直在尋找方法來加快我的研究和管理我的實驗,特別是圍繞著寫訓練管道和管理實驗配置檔案這兩個方面,我發現這兩個新專案叫做PyTorch Lightning和Hydra。PyTorch Lightning可以幫助你快速編寫訓練管道,而Hydra可以幫助你有效地管理配置檔案。

為了練習使用它們,我決定為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)的可學習引數(gammabeta)。這樣做是為了使權值檔案的格式(只有一行表示層權值,另一行表示偏差)在新增歸一化層時不必更改。

目前,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),來忽略gammabeta。所以,我把歸一化層設為兩個都忽略,然後簡單地在後面加上一個張量,它表示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/