Fix weight tests

microsoft · Feb 21, 2025 · d8e8fe6 · d8e8fe6
1 parent 1d80adc
commit d8e8fe6
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/tests/trainers/test_instance_segmentation.py b/tests/trainers/test_instance_segmentation.py
@@ -111,8 +111,8 @@ def test_invalid_backbone(self) -> None:
         with pytest.raises(ValueError, match=match):
             InstanceSegmentationTask(backbone='invalid_backbone')
 
-    def test_pretrained_backbone(self) -> None:
-        InstanceSegmentationTask(backbone='resnet50', weights=True)
+    def test_weights(self) -> None:
+        InstanceSegmentationTask(weights=True, num_classes=91)
 
     def test_no_plot_method(self, monkeypatch: MonkeyPatch, fast_dev_run: bool) -> None:
         monkeypatch.setattr(VHR10DataModule, 'plot', plot)

diff --git a/torchgeo/trainers/instance_segmentation.py b/torchgeo/trainers/instance_segmentation.py
@@ -73,7 +73,9 @@ def configure_models(self) -> None:
         weights_backbone = None
         if self.hparams['weights']:
             weights = MaskRCNN_ResNet50_FPN_Weights.COCO_V1
-            weights_backbone = ResNet50_Weights.IMAGENET1K_V1
+            # TODO: drop last layer of weights
+            if num_classes == 91:
+                weights_backbone = ResNet50_Weights.IMAGENET1K_V1
 
         # Create model
         if model == 'mask_rcnn':