Implement mixed precision for training and inference

src/membrain_seg/segmentation/segment.py

@@ -91,16 +91,18 @@ def segment(
     print("Performing 8-fold test-time augmentation.")
     for m in range(8):
         with torch.no_grad():
-            predictions += (
-                get_mirrored_img(
-                    inferer(get_mirrored_img(new_data.clone(), m).to(device), pl_model)[
-                        0
-                    ],
-                    m,
+            with torch.cuda.amp.autocast():
+                predictions += (
+                    get_mirrored_img(
+                        inferer(
+                            get_mirrored_img(new_data.clone(), m).to(device), pl_model
+                        )[0],
+                        m,
+                    )
+                    .detach()
+                    .cpu()
                 )
-                .detach()
-                .cpu()
-            )
+
     predictions /= 8.0
 
     # Extract segmentations and store them in an output file.

src/membrain_seg/segmentation/train.py

@@ -91,8 +91,7 @@ def train(
         save_top_k=-1,  # Save all checkpoints
         every_n_epochs=100,
         dirpath="checkpoints/",
-        filename=checkpointing_name
-        + "-{epoch}-{val_loss:.2f}",  # Customize the filename of saved checkpoints
+        filename=checkpointing_name + "-{epoch}-{val_loss:.2f}",
         verbose=True,  # Print a message when a checkpoint is saved
     )
 
@@ -106,6 +105,7 @@ def on_epoch_start(self, trainer, pl_module):
     print_lr_cb = PrintLearningRate()
     # Set up the trainer
     trainer = pl.Trainer(
+        precision="16-mixed",
         logger=[csv_logger, wandb_logger],
         callbacks=[
             checkpoint_callback_val_loss,

src/membrain_seg/segmentation/training/optim_utils.py

@@ -2,7 +2,10 @@
 from monai.losses import DiceLoss, MaskedLoss
 from monai.networks.nets import DynUNet
 from monai.utils import LossReduction
-from torch.nn.functional import binary_cross_entropy, sigmoid
+from torch.nn.functional import (
+    binary_cross_entropy_with_logits,
+    sigmoid,
+)
 from torch.nn.modules.loss import _Loss
 
 
@@ -79,14 +82,18 @@ def forward(self, data: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
             The calculated loss.
         """
         # Create a mask to ignore the specified label in the target
+        orig_data = data.clone()
         data = sigmoid(data)
         mask = target != self.ignore_label
 
         # Compute the cross entropy loss while ignoring the ignore_label
         target_comp = target.clone()
         target_comp[target == self.ignore_label] = 0
         target_tensor = torch.tensor(target_comp, dtype=data.dtype, device=data.device)
-        bce_loss = binary_cross_entropy(data, target_tensor, reduction="none")
+
+        bce_loss = binary_cross_entropy_with_logits(
+            orig_data, target_tensor, reduction="none"
+        )
         bce_loss[~mask] = 0.0
         bce_loss = torch.sum(bce_loss) / torch.sum(mask)
         dice_loss = self.dice_loss(data, target, mask)