implement gradient masking with cropping

test/loss/test_loss_wrapper.py

@@ -32,6 +32,61 @@ def test_masking(self):
         # print((grad[~mask] == 0).sum())
         self.assertTrue((grad[~mask] == 0).all())
 
+    def test_ApplyMask_output_shape_crop(self):
+        from torch_em.loss import ApplyMask
+
+        # _crop batch_size=1
+        shape = (1, 1, 10, 128, 128)
+        p = torch.rand(*shape)
+        t = torch.rand(*shape)
+        m = torch.rand(*shape) > .5
+        p_masked, t_masked = ApplyMask()(p, t, m)
+        out_shape = (m.sum(), shape[1])
+        self.assertTrue(p_masked.shape == out_shape)
+        self.assertTrue(t_masked.shape == out_shape)
+
+        # _crop batch_size>1
+        shape = (5, 1, 10, 128, 128)
+        p = torch.rand(*shape)
+        t = torch.rand(*shape)
+        m = torch.rand(*shape) > .5
+        p_masked, t_masked = ApplyMask()(p, t, m)
+        out_shape = (m.sum(), shape[1])
+        self.assertTrue(p_masked.shape == out_shape)
+        self.assertTrue(t_masked.shape == out_shape)
+
+        # _crop n_channels>1
+        shape = (1, 2, 10, 128, 128)
+        p = torch.rand(*shape)
+        t = torch.rand(*shape)
+        m = torch.rand(*shape) > .5
+        with self.assertRaises(ValueError):
+            p_masked, t_masked = ApplyMask()(p, t, m)
+
+        # _crop different shapes
+        shape_pt = (5, 2, 10, 128, 128)
+        p = torch.rand(*shape_pt)
+        t = torch.rand(*shape_pt)
+        shape_m = (5, 1, 10, 128, 128)
+        m = torch.rand(*shape_m) > .5
+        p_masked, t_masked = ApplyMask()(p, t, m)
+        out_shape = (m.sum(), shape_pt[1])
+        self.assertTrue(p_masked.shape == out_shape)
+        self.assertTrue(t_masked.shape == out_shape)
+
+    def test_ApplyMask_output_shape_multiply(self):
+        from torch_em.loss import ApplyMask
+
+        # _multiply
+        shape = (2, 5, 10, 128, 128)
+        p = torch.rand(*shape)
+        t = torch.rand(*shape)
+        m = torch.rand(*shape) > .5
+
+        p_masked, t_masked = ApplyMask(masking_method="multiply")(p, t, m)
+        self.assertTrue(p_masked.shape == shape)
+        self.assertTrue(t_masked.shape == shape)
+
 
 if __name__ == '__main__':
     unittest.main()

torch_em/loss/wrapper.py

@@ -1,3 +1,4 @@
+import torch
 import torch.nn as nn
 
 
@@ -43,30 +44,67 @@ def forward(self, prediction, target, **kwargs):
 
 
 class ApplyMask:
-    def __call__(self, prediction, target, mask):
-        mask.requires_grad = False
+    def _crop(prediction, target, mask, channel_dim):
+        if mask.shape[channel_dim] != 1:
+            raise ValueError(
+                "_crop only supports a mask with a singleton channel axis. \
+                Please consider using masking_method=multiply."
+            )
+        mask = mask.type(torch.bool)
+        # remove singleton axis
+        mask = mask.squeeze(channel_dim)
+        # move channel axis to end
+        prediction = prediction.moveaxis(channel_dim, -1)
+        target = target.moveaxis(channel_dim, -1)
+        # output has shape N x C
+        # correct for torch_em.loss.dice.flatten_samples
+        return prediction[mask], target[mask]
+
+    def _multiply(prediction, target, mask, channel_dim):
         prediction = prediction * mask
         target = target * mask
         return prediction, target
 
+    MASKING_FUNCS = {
+        "crop": _crop,
+        "multiply": _multiply,
+    }
+
+    def __init__(self, masking_method="crop", channel_dim=1):
+        if masking_method not in self.MASKING_FUNCS.keys():
+            raise ValueError(f"{masking_method} is not available, please use one of {list(self.MASKING_FUNCS.keys())}.")
+        self.masking_func = self.MASKING_FUNCS[masking_method]
+        self.channel_dim = channel_dim
+
+        self.init_kwargs = {
+            "masking_method": masking_method,
+            "channel_dim": channel_dim,
+        }
+
+    def __call__(self, prediction, target, mask):
+        mask.requires_grad = False
+        return self.masking_func(prediction, target, mask, self.channel_dim)
+
 
-class ApplyAndRemoveMask:
+class ApplyAndRemoveMask(ApplyMask):
     def __call__(self, prediction, target):
         assert target.dim() == prediction.dim(), f"{target.dim()}, {prediction.dim()}"
         assert target.size(1) == 2 * prediction.size(1), f"{target.size(1)}, {prediction.size(1)}"
         assert target.shape[2:] == prediction.shape[2:], f"{str(target.shape)}, {str(prediction.shape)}"
         seperating_channel = target.size(1) // 2
         mask = target[:, seperating_channel:]
         target = target[:, :seperating_channel]
-        prediction, target = ApplyMask()(prediction, target, mask)
+        prediction, target = super().__call__(prediction, target, mask)
         return prediction, target
 
 
-class MaskIgnoreLabel:
-    def __init__(self, ignore_label=-1):
+class MaskIgnoreLabel(ApplyMask):
+    def __init__(self, ignore_label=-1, masking_method="crop", channel_dim=1):
+        super().__init__(masking_method, channel_dim)
         self.ignore_label = ignore_label
+        self.init_kwargs["ignore_label"] = ignore_label
 
     def __call__(self, prediction, target):
         mask = (target != self.ignore_label)
-        prediction, target = ApplyMask()(prediction, target, mask)
+        prediction, target = super().__call__(prediction, target, mask)
         return prediction, target