From 7e3066a4a82054ceb18ee95f97124e7ae8237fd8 Mon Sep 17 00:00:00 2001 From: xmyqsh Date: Tue, 19 Oct 2021 22:29:35 +0800 Subject: [PATCH] fix freeze of bn in training --- networks/resnet_css.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/networks/resnet_css.py b/networks/resnet_css.py index dd07861..5015f59 100644 --- a/networks/resnet_css.py +++ b/networks/resnet_css.py @@ -97,6 +97,7 @@ def forward(self, x): def _freeze_module(module): + module.eval() for param in module.parameters(): param.requires_grad = False @@ -152,10 +153,15 @@ def __init__(self, block, layers, num_classes=2): nn.init.constant_(m.weight, 1) nn.init.constant_(m.bias, 0) + def freeze(self): # Freeze first layers _freeze_module(self.conv1) _freeze_module(self.bn1) _freeze_module(self.layer1) + + def train(self, mode=True): + super(ResNet, self).train(mode) + freeze() def _make_layer(self, block, planes, blocks, stride=1): downsample = None