From b3abecab415b920d23ced2a3fba143aa0d85145e Mon Sep 17 00:00:00 2001 From: Vasilis Vryniotis Date: Mon, 24 May 2021 20:02:29 +0100 Subject: [PATCH 1/2] Fix a bug when trainable_layers == 0 --- torchvision/models/detection/backbone_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchvision/models/detection/backbone_utils.py b/torchvision/models/detection/backbone_utils.py index d662de8078a..3178a81b52c 100644 --- a/torchvision/models/detection/backbone_utils.py +++ b/torchvision/models/detection/backbone_utils.py @@ -151,7 +151,7 @@ def mobilenet_backbone( # find the index of the layer from which we wont freeze assert 0 <= trainable_layers <= num_stages - freeze_before = num_stages if trainable_layers == 0 else stage_indices[num_stages - trainable_layers] + freeze_before = len(backbone) if trainable_layers == 0 else stage_indices[num_stages - trainable_layers] for b in backbone[:freeze_before]: for parameter in b.parameters(): From be4f0c4da751dfb24fcecdc0fbe2f2778cf7e61c Mon Sep 17 00:00:00 2001 From: Vasilis Vryniotis Date: Mon, 24 May 2021 20:05:51 +0100 Subject: [PATCH 2/2] Fix same issue on ssd. --- torchvision/models/detection/ssd.py | 2 +- torchvision/models/detection/ssdlite.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/torchvision/models/detection/ssd.py b/torchvision/models/detection/ssd.py index 6695167d462..b6e620dba4a 100644 --- a/torchvision/models/detection/ssd.py +++ b/torchvision/models/detection/ssd.py @@ -510,7 +510,7 @@ def _vgg_extractor(backbone_name: str, highres: bool, progress: bool, pretrained # find the index of the layer from which we wont freeze assert 0 <= trainable_layers <= num_stages - freeze_before = num_stages if trainable_layers == 0 else stage_indices[num_stages - trainable_layers] + freeze_before = len(backbone) if trainable_layers == 0 else stage_indices[num_stages - trainable_layers] for b in backbone[:freeze_before]: for parameter in b.parameters(): diff --git a/torchvision/models/detection/ssdlite.py b/torchvision/models/detection/ssdlite.py index 30e4383ed5d..26378d7038d 100644 --- a/torchvision/models/detection/ssdlite.py +++ b/torchvision/models/detection/ssdlite.py @@ -147,7 +147,7 @@ def _mobilenet_extractor(backbone_name: str, progress: bool, pretrained: bool, t # find the index of the layer from which we wont freeze assert 0 <= trainable_layers <= num_stages - freeze_before = num_stages if trainable_layers == 0 else stage_indices[num_stages - trainable_layers] + freeze_before = len(backbone) if trainable_layers == 0 else stage_indices[num_stages - trainable_layers] for b in backbone[:freeze_before]: for parameter in b.parameters():