PyTorch:train模式与eval模式的那些坑

文章目录

*
前言
1. train模式与eval模式
2. BatchNorm
3. 数学原理
4. 关于BN层的加载
结束语

前言

博主在最近开发过程中不小心被 pytorchtrain模式与 eval模式坑了一下o(*≧д≦)o!!,被坑的起因就不说了,本篇将详细介绍 train模式与 eval模式误用对模型带来的影响及 BatchNorm的数学原理【新增了一个解决 BNbug的记录】。

1. train模式与eval模式

使用过 pytorch深度学习框架的小伙伴们肯定知道,通常我们在训练模型前会加上 model.train()这行代码,或者干脆不加,而在测试模型前会加上 model.test()这行代码。
先来看看这两个模式是干嘛用的:

    def train(self: T, mode: bool = True) -> T:
        r"""Sets the module in training mode.

        This has any effect only on certain modules. See documentations of
        particular modules for details of their behaviors in training/evaluation
        mode, if they are affected, e.g. :class:Dropout, :class:BatchNorm,
        etc.

        Args:
            mode (bool): whether to set training mode () or evaluation
                         mode (). Default: .

        Returns:
            Module: self
"""
        if not isinstance(mode, bool):
            raise ValueError("training mode is expected to be boolean")
        self.training = mode
        for module in self.children():
            module.train(mode)
        return self

    def eval(self: T) -> T:
        r"""Sets the module in evaluation mode.

        This has any effect only on certain modules. See documentations of
        particular modules for details of their behaviors in training/evaluation
        mode, if they are affected, e.g. :class:Dropout, :class:BatchNorm,
        etc.

        This is equivalent with :meth:self.train(False) .

        See :ref:locally-disable-grad-doc for a comparison between
        .eval() and several similar mechanisms that may be confused with it.

        Returns:
            Module: self
"""
        return self.train(False)

根据上述的官方源码,可以得到以下信息:

eval()  将 module 设置为测试模式, 对某些模块会有影响, 比如Dropout和BatchNorm, 与 self.train(False) 等效
train(mode=True)    将 module 设置为训练模式, 对某些模块会有影响, 比如Dropout和BatchNorm

DropoutBatchNorm被宠幸的原因如下:


self.dropout = nn.Dropout(p=0.5)

Dropout层可以通过随即减少神经元的连接,能够把稠密的神经网络变成稀疏的神经网络,这样可以缓解过拟合(神经网络中神经元的连接越多,模型越复杂,模型越容易过拟合)(事实上, Dropout层表现并没有那么好)。


self.bn = nn.BatchNorm2d(num_features=128)

BatchNorm层可以对 mini-batch数据进行归一化来加速神经网络训练,加速模型的收敛速度及稳定性,除此之外,还可以缓解模型层数过多引入的梯度爆炸问题。

在训练模型时,将模型的模式设置为 train很容易理解,但是我们在测试模型时,我们需要使用所有的神经网络的神经元,这个时候就需要禁止 Dropout层发挥作用了,否则的话,模型的精度会有所降低。而测试模式下的 BatchNorm层会使用训练时的均值及方差,不再使用测试模型时输入数据的均值及方差 (稍后来解释为什么要这样)

OK,有了上述的简要介绍,我们来做个小实验,来看一下 train模式与 eval模式对模型的结果会有多大的影响。默认情况下,构建好模型之后就处于 train模式:

from torchvision.models import resnet152

if __name__ == '__main__':
    model = resnet152()
    print(model.training)

倘若我们在测试模型的时候,没有将模型设置成 eval模式下会怎样呢?我从 ImageNet数据集中选了 20张图片来进行测试模型:

PyTorch:train模式与eval模式的那些坑
先看看正常情况下的结果:
import torch
from torchvision.models import resnet152
from torch.nn import functional as F
from torchvision import transforms
from PIL import Image
import pickle
import glob
import pandas as pd

if __name__ == '__main__':
    label_info = pd.read_csv('imagenet2012_label.csv')

    transform = transforms.Compose([

        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        )])

    model = resnet152()
    ckpt = torch.load('pretrained/resnet152-394f9c45.pth')
    model.load_state_dict(ckpt)
    model.eval()

    file_list = glob.glob('imgs/*.JPEG')
    file_list = sorted(file_list)
    for file in file_list:
        img = Image.open(file)

        img = transform(img)
        img = img.unsqueeze(dim=0)

        output = model(img)
        data_softmax = F.softmax(output, dim=1).squeeze(dim=0).detach().numpy()
        index = data_softmax.argmax()

        results = label_info.loc[index, ['index', 'label', 'zh_label']].array
        print('index: {}, label: {}, zh_label: {}'.format(results[0], results[1], results[2]))

结果完全正确:

index: 162, label: beagle, zh_label: 狗
index: 101, label: tusker, zh_label: 大象
index: 484, label: catamaran, zh_label: 帆船
index: 638, label: maillot, zh_label: 泳衣
index: 475, label: car_mirror, zh_label: 反光镜
index: 644, label: matchstick, zh_label: 火柴
index: 881, label: upright, zh_label: 钢琴
index: 21, label: kite, zh_label: 鸟
index: 987, label: corn, zh_label: 玉米
index: 141, label: redshank, zh_label: 鸟
index: 335, label: fox_squirrel, zh_label: 松鼠
index: 832, label: stupa, zh_label: 皇宫
index: 834, label: suit, zh_label: 西装
index: 455, label: bottlecap, zh_label: 瓶盖
index: 847, label: tank, zh_label: 坦克
index: 248, label: Eskimo_dog, zh_label: 狗
index: 92, label: bee_eater, zh_label: 鸟
index: 959, label: carbonara, zh_label: 意大利面
index: 884, label: vault, zh_label: 拱廊
index: 0, label: tench, zh_label: 鱼

接下来将模型设置为 train模式,再次进行测试,结果如下:

index: 600, label: hook, zh_label: 钩子
index: 600, label: hook, zh_label: 钩子
index: 600, label: hook, zh_label: 钩子
index: 600, label: hook, zh_label: 钩子
index: 600, label: hook, zh_label: 钩子
index: 600, label: hook, zh_label: 钩子
index: 600, label: hook, zh_label: 钩子
index: 600, label: hook, zh_label: 钩子
index: 463, label: bucket, zh_label: 水桶
index: 600, label: hook, zh_label: 钩子
index: 463, label: bucket, zh_label: 水桶
index: 600, label: hook, zh_label: 钩子
index: 600, label: hook, zh_label: 钩子
index: 600, label: hook, zh_label: 钩子
index: 600, label: hook, zh_label: 钩子
index: 463, label: bucket, zh_label: 水桶
index: 600, label: hook, zh_label: 钩子
index: 600, label: hook, zh_label: 钩子
index: 600, label: hook, zh_label: 钩子
index: 600, label: hook, zh_label: 钩子

哦豁,发生了什么!这个结果很让人意外啊,模型输出完全错误!
ResNet152不含有 Dropout层,那引起这个结果的原因就只有一个了,那就是 BatchNorm层搞的鬼。

2. BatchNorm

pytorch中, BatchNorm定义如下:

torch.nn.BatchNorm2d(num_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True, device=None, dtype=None)

- num_features: C from an expected input of size (N, C, H, W)
- eps: a value added to the denominator for numerical stability. Default: 1e-5
- momentum: the value used for the running_mean and running_var computation. Can be set to None for cumulative moving average (i.e. simple average). Default: 0.1
- affine: a boolean value that when set to True, this module has learnable affine parameters. Default: True
- track_running_stats: a boolean value that when set to True, this module tracks the running mean and variance, and when set to False, this module does not track such statistics, and initializes statistics buffers running_mean and running_var as None. When these buffers are None, this module always uses batch statistics. in both training and eval modes. Default: True

- Input: (N, C, H, W)
- Output: (N, C, H, W)(same shape as input)

搭个简单的模型看一下:

import torch
from torch import nn

seed = 10001
torch.manual_seed(seed)

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()

        self.conv = nn.Conv2d(in_channels=3, out_channels=10, kernel_size=5, stride=1)
        self.bn = nn.BatchNorm2d(num_features=10, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True)
        self.relu = nn.ReLU()
        self.pool = nn.AdaptiveAvgPool2d(output_size=1)
        self.linear = nn.Linear(in_features=10, out_features=1)

    def forward(self, x):
        x = self.conv(x)
        var, mean = torch.var_mean(x, dim=[0, 2, 3])
        print("x's mean: {}\nx's var: {}".format(mean.detach().numpy(), var.detach().numpy()))

        x = self.bn(x)
        print('-----------------------------------------------------------------------------------')
        print("x's mean: {}\nx's var: {}".format(self.bn.running_mean.numpy(), self.bn.running_var.numpy()))

        x = self.relu(x)
        x = self.pool(x)
        x = x.view(x.size(0), -1)
        output = self.linear(x)

        return output

if __name__ == '__main__':
    model = MyModel()

    inputs = torch.randn(size=(128, 3, 32, 32))
    model(inputs)

运行一下上面的模型会发现,我们手动计算卷积后的特征的均值与方差和 BatchNorm层计算出来的均值与方差并不一致,但是能发现一些端倪,手动计算的均值与 BatchNorm层计算出来的均值相差了 10倍,这个不同点就是上述参数 momentum造成的,其默认值就是 0.1

PyTorch:train模式与eval模式的那些坑

参数 momentum的值更改为 1.0,再次运行模型,此时的卷积后的特征的均值与方差和 BatchNorm层计算出来的均值与方差完全一致:

PyTorch:train模式与eval模式的那些坑

来看下参数 affineBatchNorm中的具体作用,下图分别是 affine=Trueaffine=False

PyTorch:train模式与eval模式的那些坑

PyTorch:train模式与eval模式的那些坑

很明显, affine=TrueBatchNorm层有了可训练的参数 weightbias

最后,再来看一下一个非常重要的参数: track_running_stats
注意看上面的图中 num_batches_tracked的值,当我们将参数 track_running_stats的值设置为 TrueBatchNorm就会统计送入的数据,此时的 num_batches_tracked值为 1,也就是记录了一个 mini-batch的均值 running_mean和方差 running_var。更改下代码多计算几次:

if __name__ == '__main__':
    model = MyModel()

    inputs = torch.randn(size=(128, 3, 32, 32))
    for i in range(10):
        model(inputs)
    print('num_batches_tracked: ', model.bn.num_batches_tracked.numpy())

为了更具有说服力,我们再更改下代码,来对比一下 BatchNorm是如何统计数据的:

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()

        self.conv = nn.Conv2d(in_channels=3, out_channels=10, kernel_size=5, stride=1)
        self.bn = nn.BatchNorm2d(num_features=10, eps=1e-5, momentum=1.0, affine=True, track_running_stats=True)
        self.relu = nn.ReLU()
        self.pool = nn.AdaptiveAvgPool2d(output_size=1)
        self.linear = nn.Linear(in_features=10, out_features=1)

        self.var_data = []
        self.mean_data = []

    def forward(self, x):
        x = self.conv(x)

        var, mean = torch.var_mean(x, dim=[0, 2, 3])
        self.var_data.append(var)
        self.mean_data.append(mean)

        x = self.bn(x)
        x = self.relu(x)
        x = self.pool(x)
        x = x.view(x.size(0), -1)
        output = self.linear(x)

        return output

if __name__ == '__main__':
    model = MyModel()

    for i in range(10):
        inputs = torch.randn(size=(128, 3, 32, 32))
        model(inputs)

    var = model.var_data[-1]
    mean = model.mean_data[-1]
    print("x's mean: {}\nx's var: {}".format(mean.detach().numpy(), var.detach().numpy()))
    print('-----------------------------------------------------------------------------------')
    print("x's mean: {}\nx's var: {}".format(model.bn.running_mean.numpy(), model.bn.running_var.numpy()))

PyTorch:train模式与eval模式的那些坑
与我当初想的不太一样,我以为是历史以往所有的样本的均值与方差,其实并不是,根据实际的结果来看, BatchNorm记录的均值与方差始终是最后一个 mini-batch样本的均值与方差,即只将当前的数据进行归一化。

3. 数学原理

BatchNorm算法出自 Google的一篇论文:Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift

PyTorch:train模式与eval模式的那些坑
根据论文中的公式,可以得到 BatchNorm算法的表达式y = γ ⋅ x − E ( x ) V a r ( x ) + ϵ + β \bm y = \gamma \cdot \frac {\bm x – E(\bm x)} {\sqrt{Var(\bm x) + \epsilon}} + \beta y =γ⋅Va r (x )+ϵ​x −E (x )​+β 其中,x \bm x x是输入张量的值,ϵ \epsilon ϵ是一个较小的浮点数,以防止分母为0。
BatchNorm2d为例,平均值和方差都是相对于N、H、W三个方向进行计算和平均的,具体如下:
E ( x c ) = 1 N × H × W ∑ N , H , W x c E(\bm x_c)=\frac {1} {N \times H \times W} \sum_{N,H,W} \bm x_c E (x c ​)=N ×H ×W 1 ​N ,H ,W ∑​x c ​V a r ( x c ) = 1 N × H × W ∑ N , H , W ( x c − E ( x c ) ) 2 Var(\bm x_c)=\frac {1} {N \times H \times W} \sum_{N,H,W} \bigg(\bm x_c-E(\bm x_c)\bigg)^2 Va r (x c ​)=N ×H ×W 1 ​N ,H ,W ∑​(x c ​−E (x c ​))2 根据计算公式可以知道,统计量的输出是一个大小为C的向量。

由于在求统计量的过程中包含了 mini-batch N的平均,所以 BatchNorm又称为批次归一化方法,只改变输入 tensor的数据分布,不改变 tensor的形状。

接下来再跟着公式来看下 pytorch中的 BatchNorm2d的参数:
参数 momentum控制着指数移动平均计算E ( x ) E(\bm x)E (x )和V a r ( x ) Var(\bm x)Va r (x )时的动量, 计算公式如下:
x ^ n e w = ( 1 − α ) x ^ + α x ^ t \hat x_{new} = (1 – \alpha)\hat x + \alpha \hat x_t x ^n e w ​=(1 −α)x ^+αx ^t ​ 其中α \alpha α是动量的值,x ^ t \hat x_t x ^t ​是当前的E ( x ) E(\bm x)E (x )或V a r ( x ) Var(\bm x)Va r (x )的计算值,x ^ \hat x x ^是上一步的指数移动平均的估计值,x ^ n e w \hat x_{new}x ^n e w ​是当前的指数移动平均的估计值。
参数 affine决定了是否在归一化后做仿射变换,即是否设定β \beta β和γ \gamma γ参数, affine=True表示β \beta β和γ \gamma γ是可训练的标量参数, affine=False表示β \beta β和γ \gamma γ是固定的标量参数,即β = 0 \beta=0 β=0,γ = 1 \gamma=1 γ=1。
参数 track_running_stats决定了是否使用指数移动平均来估计当前的统计参量,默认是使用的,如果设置 track_running_stats=False,则直接使用当前统计量的计算值x ^ t \hat x_t x ^t ​来对E ( x ) E(\bm x)E (x )和V a r ( x ) Var(\bm x)Va r (x )进行估计。

仿射变换 = 线性变换 + 平移

; 4. 关于BN层的加载

今天 (2023/02/17)在量化 yolov5(v6.2)时遇到了一个量化前后浮点模型与定点模型结果对不上的问题,经调试发现,原始的yolov5预训练模型保存的是整个模型(结构),且原始模型的 BatchNorm层的 epsmomentum不是默认值 1e-50.1,而是 1e-30.03。从上述 BatchNorm的表达式可以看到, eps参数,也就是ϵ \epsilon ϵ参与了运算,目的是防止分母为 0,虽然 eps参与的是一个加法运算,但作为分母,对整体的结果影响还是很大的,尤其是这种误差在多个层运算之后会被放大,导致最终的结果对不上。
yolov5作者提供的源码中,加载预训练模型时直接用了这个 model,而并不是常规的”先构建模型,再加载预训练权重”这个流程。在量化时,需要对模型进行重构,插入量化节点,我在重构过程中使用的 BatchNorm层的参数为默认值,然后再对这个模型进行权重加载的操作。ok,流程看起来没有问题。
那么问题就来了,为什么结果对不上???

先说结论: BNepsmomentum是非训练参数,可以理解为超参,模型加载的时候不加载这俩参数,因为保存模型的 state_dict时就不含这俩参数。

做了个小实验,先看下保存时模型的 state_dict

PyTorch:train模式与eval模式的那些坑
可以很清楚的看到,模型的 state_dictBN层没有 epsmomentum这俩参数。
再看下我的代码:

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=1, bias=False)

        self.bn1 = nn.BatchNorm2d(64, eps=1e-5, momentum=0.1)
        self.relu = nn.ReLU()

        self.conv2 = nn.Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=1, bias=False)

        self.bn2 = nn.BatchNorm2d(128, eps=1e-5, momentum=0.1)

        self.conv3 = nn.Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=1, bias=False)
        self.bn3 = nn.BatchNorm2d(256, eps=1e-5, momentum=0.1)

    def forward(self, x):
        x = self.relu(self.bn1(self.conv1(x)))
        x = self.relu(self.bn2(self.conv2(x)))
        x = self.relu(self.bn3(self.conv3(x)))

        return x

if __name__ == '__main__':
    model = MyModel()

    ckpt = torch.load('mymodel.pth')
    model.load_state_dict(ckpt)
    print(model)

我先保存了 BN不是默认值得模型权重,然后再将模型的 BN层修改为默认值,然后对该模型加载刚刚保存的 BN层发生改变的模型权重,结果正如上述所说的那样, BN层加载预训练权重时 epsmomentum参数不会被修改,仍是模型构建时的初始值。

PyTorch:train模式与eval模式的那些坑
根据这次消除 Bug的经历,有必要清楚一点,那就是 BN层的 eps参数不要随意更改,或者说,在使用预训练模型时,需要留意下当前的 BN参数与原始的是否一致,否则复现出来的结果可能与原始的有出入。

结束语

在实际应用中,通常会将 mini-batch设置稍微大些,比如128, 256,如果设置的太小,可能会导致数据变化很剧烈,模型很难收敛,毕竟 mini-batch只是数据集中的很小一部分数据。

Original: https://blog.csdn.net/qq_42730750/article/details/123822902
Author: 夏小悠
Title: PyTorch:train模式与eval模式的那些坑

原创文章受到原创版权保护。转载请注明出处:https://www.johngo689.com/628331/

转载文章受原作者版权保护。转载请注明原作者出处!

(0)

大家都在看

亲爱的 Coder【最近整理,可免费获取】👉 最新必读书单  | 👏 面试题下载  | 🌎 免费的AI知识星球