add pytorch modifier

Author: seanytak

contrib/segmentation/job.yml

@@ -0,0 +1,23 @@
+$schema: https://azuremlschemas.azureedge.net/latest/commandJob.schema.json
+code:
+  local_path: .
+command: >-
+  python train.py
+  --train-dir {inputs.dummy}
+  --val-dir {inputs.dummy}
+  --patch-dim "256, 256"
+inputs:
+  dummy:
+    data:
+      local_path: data/train
+    mode: mount
+compute:
+  target: azureml:cpu-cluster
+  instance_count: 1
+# distribution:
+#   type: pytorch
+#   process_count: 2
+# azureml:<environment-name>:<version>
+environment: azureml:semantic-segmentation:1
+experiment_name: pytorch-semantic-segmentation
+description: Train a Semantic Segmentation Model on the Semantic Segmentation Drone Dataset

contrib/segmentation/src/datasets/semantic_segmentation.py

@@ -9,6 +9,7 @@
 import numpy as np
 import torch
 from PIL import Image
+from torch.utils.data.dataset import Dataset
 
 from .coco import CocoDataset
 
@@ -398,21 +399,22 @@ def __len__(self):
         return self.length
 
 
-class SemanticSegmentationDataset(torch.utils.data.Dataset):
+class SemanticSegmentationPyTorchDataset(torch.utils.data.Dataset):
 
     _available_patch_strategies = set(
         ["resize", "deterministic_center_crop", "crop_all"]
     )
 
     # NC24sv3 Azure VMs have 440GiB of RAM
-    # This allows the SemanticSegmentationDataset to be stored in memory
+    # This allows the SemanticSegmentationPyTorchDataset to be stored in memory
     # However, when multiple workers are used in PyTorch Dataloader,
     # a separate deepcopy of the dataset is made per instance
     # Thus, disk is currently the only shared memory pool between processes
     _available_cache_strategies = set([None, "none", "disk"])
 
     def __init__(
         self,
+        dataset: Dataset,
         labels_filepath: str,
         classes: List[int],
         annotation_format: str,
@@ -427,15 +429,15 @@ def __init__(
     ):
         if (
             patch_strategy
-            not in SemanticSegmentationDataset._available_patch_strategies
+            not in SemanticSegmentationPyTorchDataset._available_patch_strategies
         ):
             raise ValueError(
                 f"Parameter `patch_strategy` must be one of {self._available_patch_strategies}"
             )
 
         if (
             cache_strategy
-            not in SemanticSegmentationDataset._available_cache_strategies
+            not in SemanticSegmentationPyTorchDataset._available_cache_strategies
         ):
             raise ValueError(
                 f"Parameter `cache_strategy` must be one of {self._available_cache_strategies}"
@@ -456,24 +458,19 @@ def __init__(
                 'Parameter `patch_dim` must not be None if `patch_strategy is "crop_all"'
             )
 
-        coco = CocoDataset(
-            labels_filepath=labels_filepath,
-            root_dir=root_dir,
-            classes=classes,
-            annotation_format=annotation_format,
-        )
-
         if patch_strategy == "resize":
-            self.dataset = SemanticSegmentationResizeDataset(coco, resize_dim)
+            self.dataset = SemanticSegmentationResizeDataset(
+                dataset, resize_dim
+            )
         elif patch_strategy == "deterministic_center_crop":
             self.dataset = (
                 SemanticSegmentationWithDeterministicPatchingDataset(
-                    coco, patch_dim
+                    dataset, patch_dim
                 )
             )
         elif patch_strategy == "crop_all":
             self.dataset = SemanticSegmentationDatasetFullCoverage(
-                coco, patch_dim
+                dataset, patch_dim
             )
 
         self.root_dir = root_dir
@@ -570,7 +567,7 @@ class ToySemanticSegmentationDataset(torch.utils.data.Dataset):
     """Toy semantic segmentation dataset for integration testing purposes"""
 
     def __init__(self, *args, **kwargs):
-        self._dataset = SemanticSegmentationDataset(*args, **kwargs)
+        self._dataset = SemanticSegmentationPyTorchDataset(*args, **kwargs)
 
     def __getitem__(self, idx):
         return self._dataset[idx]

contrib/segmentation/tests/datasets/test_semantic_segmentation.py

@@ -5,7 +5,7 @@
 from PIL import Image
 
 from src.datasets.semantic_segmentation import (
-    SemanticSegmentationDataset,
+    SemanticSegmentationPyTorchDataset,
 )
 
 
@@ -53,7 +53,7 @@ def test_semantic_segmentation_dataset(
         "src.datasets.semantic_segmentation.Image.open",
         return_value=high_resolution_image,
     )
-    dataset = SemanticSegmentationDataset(
+    dataset = SemanticSegmentationPyTorchDataset(
         standard_labels_filepath,
         root_dir="data",
         classes=classes,

contrib/segmentation/train.py

@@ -16,7 +16,7 @@
 from torch.utils.data.dataloader import DataLoader
 
 from src.datasets.semantic_segmentation import (
-    SemanticSegmentationDataset,
+    SemanticSegmentationPyTorchDataset,
     SemanticSegmentationStochasticPatchingDataset,
     ToySemanticSegmentationDataset,
 )
@@ -151,7 +151,7 @@ def forward(self, x):
         default="",
     )
     parser.add_argument("--toy", type=bool, required=False, default=False)
-    parser.add_argument("--classes", type=str, default="1, 2, 3, 4")
+    parser.add_argument("--classes", type=str, default="1, 2")
     parser.add_argument(
         "--log-file", type=str, required=False, default="train.log"
     )
@@ -195,7 +195,7 @@ def forward(self, x):
         "--class-balance", type=str2bool, required=False, default=False
     )
     parser.add_argument(
-        "--cache-strategy", type=str, required=False, default="memory"
+        "--cache-strategy", type=str, required=False, default="none"
     )
     args = parser.parse_args()
 
@@ -204,8 +204,8 @@ def forward(self, x):
 
     train_dir = str(args.train_dir)
     val_dir = str(args.val_dir)
-    experiment_dir = str(uuid.uuid4())
-    model_dir = join(train_dir, experiment_dir)
+
+    model_dir = join("outputs", "models")
     Path(model_dir).mkdir(parents=True, exist_ok=True)
 
     if args.cache_dir is not None:
@@ -285,7 +285,7 @@ def forward(self, x):
     )
     # Toy Dataset for Integration Testing Purposes
     Dataset = (
-        SemanticSegmentationDataset
+        SemanticSegmentationPyTorchDataset
         if not is_toy
         else ToySemanticSegmentationDataset
     )
@@ -355,14 +355,18 @@ def forward(self, x):
         f"Validation dataset number of images: {dataset_val_len} | Batch size: {batch_size} | Expected number of batches: {tot_validation_batches}"
     )
 
-    num_classes: int = classes[-1] + 1  # Plus 1 for background
-    classes = [class_id_to_class_name[i] for i in range(num_classes)]
+    num_classes: int = len(classes) + 1  # Plus 1 for background
 
     # define training and validation data loaders
     # drop_last True to avoid single instances which throw an error on batch norm layers
 
     # Maxing the num_workers at 8 due to shared memory limitations
-    num_workers = min(int(round(multiprocessing.cpu_count() * 2 / 3)), 8)
+    num_workers = min(
+        # Preferably use 2/3's of total cpus. If the cpu count is 1, it will be set to 0 which will result
+        # in dataloader using the main thread
+        int(round(multiprocessing.cpu_count() * 2 / 3)),
+        8,
+    )
 
     dataloader = DataLoader(
         dataset,
@@ -386,7 +390,9 @@ def forward(self, x):
         model = get_fcn_resnet50(num_classes, pretrained=pretrained)
     elif model_name == "deeplab":
         model = DeepLabModelWrapper(
-            num_classes, pretrained=pretrained
+            num_classes,
+            pretrained=pretrained,
+            is_feature_extracting=pretrained,
         )  # get_deeplabv3(num_classes, is_feature_extracting=pretrained)
     else:
         raise ValueError(