关于Pytorch的BN
目前在做一些人脸方面的内容,场景是在一个训练完成的模型上加一些功能(要是模型也能支持装饰器就好了😄)。
假定现在已经有了一个训练好的模型 A \mathbf{A} A(包含了BN和Dropout层),现在在 A \mathbf{A} A的基础上额外添加了一些层 L L 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 ...