Skip to content

Commit 58a60b2

Browse files
authored
Fixing bug on SSD backbone freezing. (#4591)
# Conflicts: # torchvision/models/detection/ssd.py
1 parent ccf3aa2 commit 58a60b2

File tree

1 file changed

+2
-2
lines changed
  • torchvision/models/detection

1 file changed

+2
-2
lines changed

torchvision/models/detection/ssd.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -505,7 +505,7 @@ def _vgg_extractor(backbone_name: str, highres: bool, progress: bool, pretrained
505505
backbone = vgg.__dict__[backbone_name](pretrained=pretrained, progress=progress).features
506506

507507
# Gather the indices of maxpools. These are the locations of output blocks.
508-
stage_indices = [i for i, b in enumerate(backbone) if isinstance(b, nn.MaxPool2d)]
508+
stage_indices = [0] + [i for i, b in enumerate(backbone) if isinstance(b, nn.MaxPool2d)][:-1]
509509
num_stages = len(stage_indices)
510510

511511
# find the index of the layer from which we wont freeze
@@ -569,7 +569,7 @@ def ssd300_vgg16(pretrained: bool = False, progress: bool = True, num_classes: i
569569
warnings.warn("The size of the model is already fixed; ignoring the argument.")
570570

571571
trainable_backbone_layers = _validate_trainable_layers(
572-
pretrained or pretrained_backbone, trainable_backbone_layers, 5, 5)
572+
pretrained or pretrained_backbone, trainable_backbone_layers, 5, 4)
573573

574574
if pretrained:
575575
# no need to download the backbone if pretrained is set

0 commit comments

Comments
 (0)