From 6a3c429be8fc5198479e415fb3097295f2910c92 Mon Sep 17 00:00:00 2001 From: Vasilis Vryniotis Date: Mon, 11 Oct 2021 18:02:47 +0100 Subject: [PATCH] Fixing bug on SSD backbone freezing. # Conflicts: # torchvision/models/detection/ssd.py --- torchvision/models/detection/ssd.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/torchvision/models/detection/ssd.py b/torchvision/models/detection/ssd.py index e67c4930b30..34e64a88d1e 100644 --- a/torchvision/models/detection/ssd.py +++ b/torchvision/models/detection/ssd.py @@ -505,7 +505,7 @@ def _vgg_extractor(backbone_name: str, highres: bool, progress: bool, pretrained backbone = vgg.__dict__[backbone_name](pretrained=pretrained, progress=progress).features # 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) # 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 warnings.warn("The size of the model is already fixed; ignoring the argument.") 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: # no need to download the backbone if pretrained is set