pytorch那些坑——你确定你真的冻结了BN层?!
2020-01-09 本文已影响0人
Cat丹
最近做实例分割项目,想着直接在物体检测框架的模型上添加mask分支,冻结detection参数,只训练mask相关的参数。
for p in self.detection_net:
for param in p.parameters():
param.requires_grad = False
然而事实却是,detection相关的性能指标一直在变!简言之,没有冻结?!
打印网络层权值,发现冻结层的参数并没有改变!那么问题在哪里呢?仔细检查,发现竟然是BN层的runing_mean和runing_var在变!这两个值是统计得来的,并没有在梯度回传的轮回中。所以,param.requires_grad=False对它们不起任何作用!
那么,正确打开方式是什么呢?直接使用eval模式。
def fix_bn(m):
classname = m.__class__.__name__
if classname.find('BatchNorm') != -1:
m.eval()
model = models.resnet50(pretrained=True)
model.cuda()
model.train()
model.apply(fix_bn) # fix batchnorm