PyTorch中如何将Norm层的running stats重置为默认初始值

需重置的BatchNorm层参数 #

你需要重置BatchNorm层的三个运行时统计参数，均为无梯度的持久化缓冲区参数：

• running_mean：默认初始值为全0张量，维度和BN层通道数一致
• running_var：默认初始值为全1张量，维度和BN层通道数一致
• num_batches_tracked：默认初始值为0，用于统计训练步数计算滑动平均的动量

补全后的代码示例 #

import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv0 = nn.Conv2d(3, 3, 3, padding = 1)
        self.norm = nn.BatchNorm2d(3)
        self.conv = nn.Conv2d(3, 3, 3, padding = 1)
    
    def forward(self, x):
        x = self.conv0(x)
        x = self.norm(x)
        return self.conv(x)

net = Net()
# 加载预训练模型的场景替换为下面这行
# net = torch.load('net.pth')

def drop_to_default():
    for m in net.modules():
        # 用isinstance兼容继承BatchNorm2d的自定义层，比type判断适用性更广
        if isinstance(m, nn.BatchNorm2d):
            # 重置running_mean为全0
            nn.init.constant_(m.running_mean, 0.)
            # 重置running_var为全1
            nn.init.constant_(m.running_var, 1.)
            # 兼容旧版本PyTorch无num_batches_tracked参数的情况
            if hasattr(m, 'num_batches_tracked'):
                nn.init.constant_(m.num_batches_tracked, 0)

drop_to_default()

可选补充说明 #

如果你的场景不需要保留预训练得到的BN层缩放参数weight（gamma）和偏移参数bias（beta），可以在上述if判断块内补充对这两个可训练参数的重置：

nn.init.constant_(m.weight, 1.)
nn.init.constant_(m.bias, 0.)

如果仅需要重新统计新数据集的滑动均值和方差、保留原有gamma和beta，就不需要添加上述两行代码。

内容的提问来源于stack exchange，提问作者XMaSt3R