用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/