torch冻结训练中bn问题

今天帮同事修改模型冻结训练中bug，记录一下
模型冻结训练中我们采用的冻结方式为

  freeze = ['', ]  # parameter names to freeze (full or partial)
    if any(freeze):
        for k, v in model.named_parameters():
            if any(x in k for x in freeze):
                print('freezing %s' % k)
                v.requires_grad = False

 在对比输出结果的时候发现精度有所降低，遂通过排查发现模型中batchnormal模块中run_meaning的值发生变化，该参数并未冻结。因此针对该问题经过查阅相关资料
# 冻结bn层的running_mean 和running_var

def freeze_bn(m):
    # freeze bn(running_mean and running_var), because pytorch (requires_grad = False) just freeze bn.weights and bn.bias
    for i,k in m.named_children():
        # print(k)
        if isinstance(k,nn.BatchNorm2d):
            # print(k.__class__.__name__)
            k.eval()
        else:
            freeze_bn(k)

其中m为深度学习模型，freeze_bn中采用递归的方法主要是由于model采用多个模块进行组装