torch_em.loss.wrapper

  1from typing import Callable, Sequence, Tuple, Union
  2
  3import torch
  4import torch.nn as nn
  5
  6
  7class LossWrapper(nn.Module):
  8    """A wrapper around a torch loss function.
  9
 10    Applies transformations to prediction and/or target before passing it to the loss.
 11
 12    Args:
 13        loss: The loss function.
 14        transform: The transformation applied to prediction and/or target.
 15            Must take both the prediction and target as arguments and then return them, potentially transformed.
 16    """
 17    def __init__(self, loss: nn.Module, transform: Callable):
 18        super().__init__()
 19        self.loss = loss
 20
 21        if not callable(transform):
 22            raise ValueError("transform has to be callable.")
 23        self.transform = transform
 24        self.init_kwargs = {'loss': loss, 'transform': transform}
 25
 26    def apply_transform(self, prediction, target, **kwargs):
 27        """@private
 28        """
 29        # Check if the prediction and target are lists.
 30        # If they are, apply the transform to each element individually.
 31        if isinstance(prediction, (list, tuple)):
 32            assert isinstance(target, (list, tuple))
 33            transformed_prediction, transformed_target = [], []
 34            for pred, targ in zip(prediction, target):
 35                tr_pred, tr_targ = self.transform(pred, targ, **kwargs)
 36                transformed_prediction.append(tr_pred)
 37                transformed_target.append(tr_targ)
 38            return transformed_prediction, transformed_target
 39        # Otherwise, we expect that prediction and target are both tensors.
 40        else:
 41            prediction, target = self.transform(prediction, target, **kwargs)
 42            return prediction, target
 43
 44    def forward(
 45        self,
 46        prediction: Union[Sequence[torch.Tensor], torch.Tensor],
 47        target: Union[Sequence[torch.Tensor], torch.Tensor],
 48        **kwargs
 49    ) -> torch.Tensor:
 50        """Apply the tranformations to prediction and/or target before computing the loss.
 51
 52        Args:
 53            prediction: The prediction.
 54            target: The target.
 55            kwargs: Additional keyword arguments for the transformation.
 56
 57        Returns:
 58            The loss.
 59        """
 60        prediction, target = self.apply_transform(prediction, target, **kwargs)
 61        loss = self.loss(prediction, target)
 62        return loss
 63
 64
 65#
 66# Loss transformations
 67#
 68def _crop(prediction, target, mask, channel_dim):
 69    if mask.shape[channel_dim] != 1:
 70        raise ValueError(
 71            "_crop only supports a mask with a singleton channel axis. Please consider using masking_method=multiply."
 72        )
 73    mask = mask.type(torch.bool)
 74    # remove singleton axis
 75    mask = mask.squeeze(channel_dim)
 76    # move channel axis to end
 77    prediction = prediction.moveaxis(channel_dim, -1)
 78    target = target.moveaxis(channel_dim, -1)
 79    # output has shape N x C
 80    # correct for torch_em.loss.dice.flatten_samples
 81    return prediction[mask], target[mask]
 82
 83
 84def _multiply(prediction, target, mask, channel_dim):
 85    prediction = prediction * mask
 86    target = target * mask
 87    return prediction, target
 88
 89
 90class ApplyMask:
 91    """Apply a mask to prediction and target, so that only values in the mask are taken into account for the loss.
 92
 93    Supports two different masking methods:
 94    - 'crop': Crop away the mask from the prediction and target. This only works if the mask just has a single channel,
 95        and if the loss function does not require spatial inputs.
 96    - 'multiply': Multiply the prediction and target with zeros outside of the mask.
 97
 98    Args:
 99        masking_method: The masking method to use. Can be one of 'crop' or 'multiply'.
100        channel_dim: The dimension of the channel axis.
101    """
102    MASKING_FUNCS = {"crop": _crop, "multiply": _multiply}
103
104    def __init__(self, masking_method: str = "crop", channel_dim: int = 1):
105        if masking_method not in self.MASKING_FUNCS.keys():
106            raise ValueError(f"{masking_method} is not available, please use one of {list(self.MASKING_FUNCS.keys())}.")
107        self.masking_func = self.MASKING_FUNCS[masking_method]
108        self.channel_dim = channel_dim
109        self.init_kwargs = {"masking_method": masking_method, "channel_dim": channel_dim}
110
111    def __call__(
112        self, prediction: torch.Tensor, target: torch.Tensor, mask: torch.Tensor
113    ) -> Tuple[torch.Tensor, torch.Tensor]:
114        """Mask predictions.
115
116        Args:
117            prediction: The prediction tensor.
118            target: The target tensor.
119            mask: The mask tensor.
120
121        Returns:
122            The masked prediction.
123            The masked target.
124        """
125        mask.requires_grad = False
126        return self.masking_func(prediction, target, mask, self.channel_dim)
127
128
129class ApplyAndRemoveMask(ApplyMask):
130    """Extract mask from extra channels from a target tensor and use it to mask the prediction.
131
132    Supports the same masking methods as `ApplyMask`.
133    """
134    def __call__(self, prediction: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
135        """Remove masking channels from the target and then apply the mask.
136
137        Args:
138            prediction: The prediction tensor.
139            target: The target tensor, with extra channels that contain the mask.
140
141        Returns:
142            The masked prediction.
143            The masked target, with extra channels removed.
144        """
145        assert target.dim() == prediction.dim(), f"{target.dim()}, {prediction.dim()}"
146        assert target.size(1) == 2 * prediction.size(1), f"{target.size(1)}, {prediction.size(1)}"
147        assert target.shape[2:] == prediction.shape[2:], f"{str(target.shape)}, {str(prediction.shape)}"
148        seperating_channel = target.size(1) // 2
149        mask = target[:, seperating_channel:]
150        target = target[:, :seperating_channel]
151        prediction, target = super().__call__(prediction, target, mask)
152        return prediction, target
153
154
155class MaskIgnoreLabel(ApplyMask):
156    """Mask ignore label from the target.
157
158    Supports the same masking methods as `ApplyMask`.
159
160    Args:
161        ignore_label: The ignore label, which will be msaked.
162        masking_method: The masking method to use. Can be one of 'crop' or 'multiply'.
163        channel_dim: The dimension of the channel axis.
164    """
165    def __init__(self, ignore_label: int = -1, masking_method: str = "crop", channel_dim: int = 1):
166        super().__init__(masking_method, channel_dim)
167        self.ignore_label = ignore_label
168        self.init_kwargs["ignore_label"] = ignore_label
169
170    def __call__(self, prediction: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
171        """Mask ignore label in the prediction and target.
172
173        Args:
174            prediction: The prediction tensor.
175            target: The target tensor.
176
177        Returns:
178            The masked prediction.
179            The masked target.
180        """
181        mask = (target != self.ignore_label)
182        prediction, target = super().__call__(prediction, target, mask)
183        return prediction, target