折腾MNIST系列——增加LayerNorm减少训练时间

前言

本文在官方MNIST代码基础上，增加LayerNorm层，从实验结果看， LayerNorm的确如原论文所说，加快了网络训练速度。

1. 动手做

1.1 基础模型

两层线性网络+Relu+Softmax+交叉熵

class MLPNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(28*28, 512)
        self.fc2 = nn.Linear(512, 10)
        
    def forward(self, x):
        out = x
        in_size = x.size(0)
        out = out.view(in_size, -1)
        out = self.fc1(out) # 1 * 512
        out = F.relu(out)
        out = self.fc2(out) # 1 * 10
        out = F.softmax(out, dim=1)
        return out

1.2 增加LayerNorm

1.2.1 LayerNorm默认设置

class MLPNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(28*28, 512)
        self.fc2 = nn.Linear(512, 10)
        self.lm = nn.LayerNorm(512)
        
    def forward(self, x):
        out = x
        in_size = x.size(0)
        out = out.view(in_size, -1)
        out = self.fc1(out) # 1 * 512
        out = self.lm(out) # 1 * 512
        out = F.relu(out)
        out = self.fc2(out) # 1 * 10
        out = F.softmax(out, dim=1)
        return out

1.2.2 LayerNorm自定义

去掉layernorm默认的仿射变换。

self.lm = nn.LayerNorm(512, elementwise_affine=False)

1.3 实验结果

20轮准确率对比

模型	初始准确率	最高准确率
基础模型	93.33%	98.18%
LayerNorm默认	94.88%	98.20%
LayerNorm自定义	94.28%	98.25%

2. 学知识

2.1 affine transform（仿射变换）

仿射变换就是：线性变换+平移。

$$Y=WX+b$$

2.2 element-wise affine

就是逐个元素应用仿射变换。

2.3 LayerNorm的实际作用

层规范化在稳定递归网络中的隐藏状态动态方面非常有效。从经验上看，与之前发表的技术相比，层规范化可以显著减少训练时间。
——原论文《Layer Normalization》

如论文所说，LayerNorm能加快收敛速度，从上图可以发现，加了LayerNorm之后，可以在10轮左右达到原始网络20轮的训练效果，还是非常给力的，毕竟这代表着节约时间和金钱。

3. 小结

LayerNorm可以显著减少训练时间，是很好的设计。
但是LayerNorm并不会提高准确率。