feat: add mixed precision inference

Author: okunator

cellseg_models_pytorch/inference/_base_inferer.py

@@ -203,19 +203,32 @@ def from_yaml(cls, model: nn.Module, yaml_path: str):
     def _infer_batch(self):
         raise NotImplementedError
 
-    def infer(self) -> None:
+    def infer(self, mixed_precision: bool = False) -> None:
         """Run inference and post-processing for the images.
 
         NOTE:
-        - Saves outputs in `self.out_masks` or to disk (.mat/.json) files.
-        - If `save_intermediate` is set to True, also intermiediate model outputs are
-            saved to `self.soft_masks`
-        - `self.out_masks` and `self.soft_masks` are nested dicts: E.g.
-                {"sample1": {"inst": [H, W], "type": [H, W], "sem": [H, W]}}
-        - If masks are saved to geojson .json files, more key word arguments
+        - Saves outputs in class attributes or to disk (.mat/.json) files.
+        - If masks are saved to .json (geojson) files, more key word arguments
             need to be given at class initialization. Namely: `geo_format`,
             `classes_type`, `classes_sem`, `offsets`. See more in the
             `FileHandler.save_masks` docs.
+
+        Attributes
+        ----------
+            - out_masks : Dict[str, Dict[str, np.ndarray]]
+                The output masks for each image. The keys are the image names and the
+                values are dictionaries of the masks. E.g.
+                {"sample1": {"inst": [H, W], "type": [H, W], "sem": [H, W]}}
+            - soft_masks : Dict[str, Dict[str, np.ndarray]]
+                NOTE: This attribute is set only if `save_intermediate = True`.
+                The soft masks for each image. I.e. the soft predictions of the trained
+                model The keys are the image names and the values are dictionaries of
+                the soft masks. E.g. {"sample1": {"type": [H, W], "aux": [C, H, W]}}
+
+        Parameters
+        ----------
+            mixed_precision : bool, default=False
+                If True, inference is performed with mixed precision.
         """
         self.soft_masks = {}
         self.out_masks = {}
@@ -227,7 +240,7 @@ def infer(self) -> None:
                     names = data["file"]
                     loader.set_description("Running inference")
                     loader.set_postfix_str("Forward pass")
-                    soft_masks = self._infer_batch(data["im"])
+                    soft_masks = self._infer_batch(data["im"], mixed_precision)
                     loader.set_postfix_str("post-processing")
                     soft_masks = self._prepare_mask_list(names, soft_masks)
 

cellseg_models_pytorch/inference/predictor.py

@@ -108,7 +108,10 @@ def __init__(
         )
 
     def forward_pass(
-        self, patch: Union[np.ndarray, torch.Tensor], in_dim_format: str = "HWC"
+        self,
+        patch: Union[np.ndarray, torch.Tensor],
+        in_dim_format: str = "HWC",
+        mixed_precision: bool = False,
     ) -> Dict[str, torch.Tensor]:
         """Input an image patch or batch of patches to the network and return logits.
 
@@ -119,7 +122,8 @@ def forward_pass(
             in_dim_format : str, default="HWC"
                 The order of the dimensions in the input array.
                 One of: "HWC", "BHWC"
-
+            mixed_precision : bool, default=False
+                Use mixed precision for inference.
 
         Returns
         -------
@@ -145,7 +149,11 @@ def forward_pass(
             patch = patch.float()
 
         with torch.no_grad():
-            out = self.model(patch)
+            if mixed_precision:
+                with torch.autocast(device_type="cuda", dtype=torch.float16):
+                    out = self.model(patch)
+            else:
+                out = self.model(patch)
 
         return out
 

cellseg_models_pytorch/inference/resize_inferer.py

@@ -139,7 +139,9 @@ def __init__(
             **kwargs,
         )
 
-    def _infer_batch(self, input_batch: torch.Tensor) -> Dict[str, torch.Tensor]:
+    def _infer_batch(
+        self, input_batch: torch.Tensor, mixed_precision: bool = False
+    ) -> Dict[str, torch.Tensor]:
         """Infer one batch of images."""
         inp_shape = tuple(input_batch.shape[2:])
 
@@ -154,7 +156,7 @@ def _infer_batch(self, input_batch: torch.Tensor) -> Dict[str, torch.Tensor]:
             input_batch = F.interpolate(input_batch, self.patch_size)
 
         batch = input_batch.to(self.device).float()
-        logits = self.predictor.forward_pass(batch)
+        logits = self.predictor.forward_pass(batch, mixed_precision=mixed_precision)
 
         probs = {}
         for k, logit in logits.items():

cellseg_models_pytorch/inference/sliding_window_inferer.py

@@ -193,7 +193,9 @@ def _get_slices(
 
         return xyslices, pady, padx
 
-    def _infer_batch(self, input_batch: torch.Tensor) -> Dict[str, torch.Tensor]:
+    def _infer_batch(
+        self, input_batch: torch.Tensor, mixed_precision: bool = False
+    ) -> Dict[str, torch.Tensor]:
         """Infer one batch of images."""
         slices, pady, padx = self._get_slices(
             self.stride, self.patch_size, tuple(input_batch.shape[2:]), self.padding
@@ -235,7 +237,7 @@ def _infer_batch(self, input_batch: torch.Tensor) -> Dict[str, torch.Tensor]:
         # run inference with the slices
         for k, (yslice, xslice) in slices.items():
             batch = input_batch[..., yslice, xslice].to(self.device).float()
-            logits = self.predictor.forward_pass(batch)
+            logits = self.predictor.forward_pass(batch, mixed_precision=mixed_precision)
 
             probs = {}
             for k, logit in logits.items():