torch_em.loss.contrastive

  1from warnings import warn
  2from typing import Optional
  3
  4import torch
  5import torch.nn as nn
  6from . import contrastive_impl as impl
  7
  8
  9def check_consecutive(labels: torch.Tensor) -> bool:
 10    """Check that the input labels are consecutive and start at zero.
 11
 12    Args:
 13        labels: The labels to check.
 14
 15    Returns:
 16        Whether the labels are consecutive.
 17    """
 18    diff = labels[1:] - labels[:-1]
 19    return (labels[0] == 0) and (diff == 1).all()
 20
 21
 22# TODO support more sophisticated ignore labels:
 23# - ignore_dist: ignored in distance term
 24# - ignore_var: ignored in variance term
 25class ContrastiveLoss(nn.Module):
 26    """Implementation of a contrastive segmentation loss.
 27
 28    From "Semantic Instance Segmentation with a Discriminative Loss Function":
 29    https://arxiv.org/pdf/1708.02551.pdf
 30
 31    This class contains different implementations for the discrimnative loss:
 32    - Based on pure pytorch, expanding the instance dimension, this is not memory efficient.
 33    - Based on pytorch_scatter (https://github.com/rusty1s/pytorch_scatter), this is memory efficient.
 34
 35    Args:
 36        delta_var: The hinge distance for the variance term.
 37            The variance term corresponds to the attractive term of the loss function.
 38        delta_dist: The hinge distance for the distance term.
 39            The distance term corresponds to the repulsive term of the loss function.
 40        norm: The norm to use.
 41        alpha: Weight for the variance term of the loss.
 42        beta: Weight for the distance term of the loss.
 43        gamma: Weight for the regularization term of the loss.
 44        ignore_label: Ignore label to exclude from the loss computation.
 45        impl: Implementation of the loss to use, either 'scatter' or 'expand'.
 46    """
 47    implementations = (None, "scatter", "expand")
 48
 49    def __init__(
 50        self,
 51        delta_var: float,
 52        delta_dist: float,
 53        norm: str = "fro",
 54        alpha: float = 1.0,
 55        beta: float = 1.0,
 56        gamma: float = 0.001,
 57        ignore_label: Optional[int] = None,
 58        impl: Optional[str] = None
 59    ):
 60        assert ignore_label is None, "Not implemented"
 61        super().__init__()
 62        self.delta_var = delta_var
 63        self.delta_dist = delta_dist
 64        self.norm = norm
 65        self.alpha = alpha
 66        self.beta = beta
 67        self.gamma = gamma
 68        self.ignore_label = ignore_label
 69
 70        assert impl in self.implementations
 71        has_torch_scatter = self.has_torch_scatter()
 72        if impl is None:
 73            if not has_torch_scatter:
 74                pt_scatter = "https://github.com/rusty1s/pytorch_scatter"
 75                warn(f"ContrastiveLoss: using pure pytorch implementation. Install {pt_scatter} for memory efficiency.")
 76            self._contrastive_impl = self._scatter_impl_batch if has_torch_scatter else self._expand_impl_batch
 77        elif impl == "scatter":
 78            assert has_torch_scatter
 79            self._contrastive_impl = self._scatter_impl_batch
 80        elif impl == "expand":
 81            self._contrastive_impl = self._expand_impl_batch
 82
 83        # all torch_em classes should store init kwargs to easily recreate the init call
 84        self.init_kwargs = {"delta_var": delta_var, "delta_dist": delta_dist, "norm": norm,
 85                            "alpha": alpha, "beta": beta, "gamma": gamma, "ignore_label": ignore_label,
 86                            "impl": impl}
 87
 88    @staticmethod
 89    def has_torch_scatter():
 90        """@private
 91        """
 92        try:
 93            import torch_scatter
 94        except ImportError:
 95            torch_scatter = None
 96        return torch_scatter is not None
 97
 98    # This implementation expands all tensors to match the instance dimensions.
 99    # Hence it's fast, but has high memory consumption.
100    # The implementation does not support masking any instance labels in the loss.
101    def _expand_impl_batch(self, input_batch, target_batch, ndim):
102        # add singleton batch dimension required for further computation
103        input_batch = input_batch.unsqueeze(0)
104
105        # get number of instances in the batch
106        instances = torch.unique(target_batch)
107        assert check_consecutive(instances), f"{instances}"
108        n_instances = instances.size()[0]
109
110        # SPATIAL = D X H X W in 3d case, H X W in 2d case
111        # expand each label as a one-hot vector: N x SPATIAL -> N x C x SPATIAL
112        target_batch = impl.expand_as_one_hot(target_batch, n_instances)
113
114        cluster_means, embeddings_per_instance = impl._compute_cluster_means(input_batch,
115                                                                             target_batch, ndim)
116        variance_term = impl._compute_variance_term(cluster_means, embeddings_per_instance,
117                                                    target_batch, ndim, self.norm, self.delta_var)
118        distance_term = impl._compute_distance_term(cluster_means, n_instances,
119                                                    ndim, self.norm, self.delta_dist)
120        regularization_term = impl._compute_regularizer_term(cluster_means, n_instances,
121                                                             ndim, self.norm)
122        # compute total loss
123        return self.alpha * variance_term + self.beta * distance_term + self.gamma * regularization_term
124
125    def _scatter_impl_batch(self, input_batch, target_batch, ndim):
126        # add singleton batch dimension required for further computation
127        input_batch = input_batch.unsqueeze(0)
128
129        instance_ids, instance_sizes = torch.unique(target_batch, return_counts=True)
130        n_instances = len(instance_ids)
131        cluster_means = impl._compute_cluster_means_scatter(input_batch, target_batch, ndim, n_lbl=n_instances)
132
133        variance_term = impl._compute_variance_term_scatter(cluster_means, input_batch, target_batch, self.norm,
134                                                            self.delta_var, instance_sizes)
135        distance_term = impl._compute_distance_term_scatter(cluster_means, self.norm, self.delta_dist)
136
137        regularization_term = torch.sum(
138            torch.norm(cluster_means, p=self.norm, dim=1)
139        ).div(n_instances)
140
141        # Compute the combined loss.
142        return self.alpha * variance_term + self.beta * distance_term + self.gamma * regularization_term
143
144    def forward(self, input_: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
145        """Compute the discrimantive loss.
146
147        Args:
148            input_: The embedding predictions.
149            target: The target segmentation.
150
151        Returns:
152            The discriminative loss value.
153        """
154        n_batches = input_.shape[0]
155        assert target.dim() == input_.dim()
156        assert target.shape[1] == 1
157        assert n_batches == target.shape[0]
158        assert input_.size()[2:] == target.size()[2:]
159
160        ndim = input_.dim() - 2
161        assert ndim in (2, 3)
162
163        # iterate over the batches
164        loss = 0.0
165        for input_batch, target_batch in zip(input_, target):
166            loss_batch = self._contrastive_impl(input_batch, target_batch, ndim)
167            loss += loss_batch
168
169        return loss.div(n_batches)