fix load from pretrained and add resnet 101

flash/vision/segmentation/model.py

@@ -63,7 +63,7 @@ def __init__(
         num_classes: int,
         backbone: Union[str, Tuple[nn.Module, int]] = "torchvision/fcn_resnet50",
         backbone_kwargs: Optional[Dict] = None,
-        pretrained: bool = False,
+        pretrained: bool = True,
         loss_fn: Optional[Callable] = None,
         optimizer: Type[torch.optim.Optimizer] = torch.optim.Adam,
         metrics: Optional[Union[Callable, Mapping, Sequence, None]] = None,
@@ -104,6 +104,14 @@ def forward(self, x) -> torch.Tensor:
 
 
 @SemanticSegmentation.backbones(name="torchvision/fcn_resnet50")
-def fn(pretrained: bool, num_classes: int) -> nn.Module:
-    model: nn.Module = torchvision.models.segmentation.fcn_resnet50(pretrained=pretrained, num_classes=num_classes)
+def load_torchvision_fcn_resnet50(pretrained: bool, num_classes: int) -> nn.Module:
+    model = torchvision.models.segmentation.fcn_resnet50(pretrained=pretrained)
+    model.classifier[-1] = nn.Conv2d(512, num_classes, kernel_size=(1, 1), stride=(1, 1))
+    return model
+
+
+@SemanticSegmentation.backbones(name="torchvision/fcn_resnet101")
+def load_torchvision_fcn_resnet101(pretrained: bool, num_classes: int) -> nn.Module:
+    model = torchvision.models.segmentation.fcn_resnet101(pretrained=pretrained)
+    model.classifier[-1] = nn.Conv2d(512, num_classes, kernel_size=(1, 1), stride=(1, 1))
     return model

flash_examples/finetuning/semantic_segmentation.py

@@ -53,13 +53,14 @@ def load_data(data_root: str = 'data/') -> Tuple[List[str], List[str]]:
     batch_size=4,
     val_split=0.3,  # TODO: this needs to be implemented
     image_size=(300, 400),  # (600, 800)
+    num_workers=0,
 )
 
 # 2.2 Visualise the samples
 labels_map = SegmentationLabels.create_random_labels_map(num_classes=21)
 datamodule.set_labels_map(labels_map)
-datamodule.show_train_batch("load_sample")
-datamodule.show_train_batch("post_tensor_transform")
+#datamodule.show_train_batch("load_sample")
+#datamodule.show_train_batch("post_tensor_transform")
 
 # 3. Build the model
 model = SemanticSegmentation(
@@ -69,9 +70,9 @@ def load_data(data_root: str = 'data/') -> Tuple[List[str], List[str]]:
 
 # 4. Create the trainer.
 trainer = flash.Trainer(
-    max_epochs=5,
+    max_epochs=1,
     gpus=1,
-    # precision=16,  # why slower ? :)
+    #precision=16,  # why slower ? :)
 )
 
 # 5. Train the model

tests/vision/segmentation/test_model.py

@@ -48,6 +48,7 @@ def test_forward(num_classes, img_shape):
     "backbone",
     [
         "torchvision/fcn_resnet50",
+        "torchvision/fcn_resnet101",
     ],
 )
 def test_init_train(tmpdir, backbone):