pytorch冻结bn

在PyTorch中，冻结Batch Normalization（BN）层的常见做法是将其设置为eval模式，并将其track_running_stats属性设置为False。这样做可以防止BN层参与训练过程中的梯度更新。以下是几种常见的冻结BN层的方法：

方法一：
在加载预训练模型时，需要使用以下代码来冻结BN层：

def freeze_bn(m):
    classname = m.__class__.__name__
    if classname.find('BatchNorm') != -1:
        m.eval()

model.apply(freeze_bn)

这段代码会将模型中所有的BN层设置为eval模式，从而冻结它们的参数。

方法二：
如果在自己定义的模型中需要冻结特征提取层（pretrain layer）的BN层，可以按如下方式修改train函数：

def train(self, mode=True):
    super(fintuneNet, self).train(mode)
    if self.args.freeze_bn and mode==True:
        self.branch_cnn.apply(self.fix_bn)
    return self

def fix_bn(self, m):
    classname = m.__class__.__name__
    if classname.find('BatchNorm') != -1:
        m.eval()
        m.track_running_stats = False
    for name, p in m.named_parameters():
        p.requires_grad = False

这段代码会将模型中特征提取层的BN层设置为eval模式，并将其track_running_stats属性设置为False，同时将参数的requires_grad属性设置为False，从而冻结这些层的参数。

另外，可以阅读一篇名为"Pytorch BN(BatchNormal)计算过程与源码分析和train与eval的区别"的文章，该文章对PyTorch中BN层的计算过程以及train和eval模式的区别进行了详细分析。