Skip to content

GitLab

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

Fixing bug on SSD backbone freezing. (#4590)

parent 221dd5f7
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
torchvision/models/detection/ssd.py
View file @ 7770b04e
...@@ -532,7 +532,7 @@ def _vgg_extractor(backbone_name: str, highres: bool, progress: bool, pretrained ...@@ -532,7 +532,7 @@ def _vgg_extractor(backbone_name: str, highres: bool, progress: bool, pretrained
backbone = vgg.__dict__[backbone_name](pretrained=pretrained, progress=progress).features backbone = vgg.__dict__[backbone_name](pretrained=pretrained, progress=progress).features
# Gather the indices of maxpools. These are the locations of output blocks. # Gather the indices of maxpools. These are the locations of output blocks.
stage_indices = [i for i, b in enumerate(backbone) if isinstance(b, nn.MaxPool2d)] stage_indices = [0] + [i for i, b in enumerate(backbone) if isinstance(b, nn.MaxPool2d)][:-1]
num_stages = len(stage_indices) num_stages = len(stage_indices)
# find the index of the layer from which we wont freeze # find the index of the layer from which we wont freeze
...@@ -602,7 +602,7 @@ def ssd300_vgg16( ...@@ -602,7 +602,7 @@ def ssd300_vgg16(
warnings.warn("The size of the model is already fixed; ignoring the argument.") warnings.warn("The size of the model is already fixed; ignoring the argument.")
trainable_backbone_layers = _validate_trainable_layers( trainable_backbone_layers = _validate_trainable_layers(
pretrained or pretrained_backbone, trainable_backbone_layers, 5, 5 pretrained or pretrained_backbone, trainable_backbone_layers, 5, 4
) )
if pretrained: if pretrained:
......
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