关于pytorch的BN，在训练的模型上增添新模块[只训练新模块]

关于Pytorch的BN

目前在做一些人脸方面的内容，场景是在一个训练完成的模型上加一些功能（要是模型也能支持装饰器就好了😄）。

假定现在已经有了一个训练好的模型$\mathbf{A}$（包含了BN和Dropout层），现在在$\mathbf{A}$的基础上额外添加了一些层$L$，需要对这些新添加的层进行训练。

那么"常规"的做法就如下：

model = build_new_model() # add requires_grad=False to those untrained parameters
model.load_state_dict(torch.load('model_A.pth'), strict=False)
optimizer = torch.optim.Adam(filter(lambda x: x.requires_grad, model.parameters()), lr=lr)
# add some codes about learning rate schedule ... 
# train part
model.train()
# some training code balabala ... 
model.eval()
# some validating/testing code balabla ...

但是这么出来的结果却不对，原因就是在于BN和Dropout层上，详细的分析可以参考这个blog。

原因在于无论你怎么设置requires_grad跟BN都没有关系，BN的moving average的参数只跟model.train()或者model.eval()有关…

纠正方法也不难，手动把BN类全部手动拉成eval模式就行。

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

model.train()
model.apply(fix_bn)
# some training code balabala ... 
model.eval()
model.apply(fix_bn)
# some validating/testing code balabla ...