Unverified Commit eca37cf7 authored by Vasilis Vryniotis's avatar Vasilis Vryniotis Committed by GitHub
Browse files

Fix incorrectly frozen BN on ResNet FPN backbone (#3396)

* Avoid freezing bn1 if all layers are trainable.

* Remove misleading comments.
parent 067b9dca
...@@ -96,7 +96,8 @@ def resnet_fpn_backbone( ...@@ -96,7 +96,8 @@ def resnet_fpn_backbone(
# select layers that wont be frozen # select layers that wont be frozen
assert 0 <= trainable_layers <= 5 assert 0 <= trainable_layers <= 5
layers_to_train = ['layer4', 'layer3', 'layer2', 'layer1', 'conv1'][:trainable_layers] layers_to_train = ['layer4', 'layer3', 'layer2', 'layer1', 'conv1'][:trainable_layers]
# freeze layers only if pretrained backbone is used if trainable_layers == 5:
layers_to_train.append('bn1')
for name, parameter in backbone.named_parameters(): for name, parameter in backbone.named_parameters():
if all([not name.startswith(layer) for layer in layers_to_train]): if all([not name.startswith(layer) for layer in layers_to_train]):
parameter.requires_grad_(False) parameter.requires_grad_(False)
...@@ -152,7 +153,6 @@ def mobilenet_backbone( ...@@ -152,7 +153,6 @@ def mobilenet_backbone(
assert 0 <= trainable_layers <= num_stages assert 0 <= trainable_layers <= num_stages
freeze_before = num_stages if trainable_layers == 0 else stage_indices[num_stages - trainable_layers] freeze_before = num_stages if trainable_layers == 0 else stage_indices[num_stages - trainable_layers]
# freeze layers only if pretrained backbone is used
for b in backbone[:freeze_before]: for b in backbone[:freeze_before]:
for parameter in b.parameters(): for parameter in b.parameters():
parameter.requires_grad_(False) parameter.requires_grad_(False)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment