torch_em.self_training.loss

  1from typing import Optional, Union
  2
  3import torch
  4import torch.nn as nn
  5
  6import torch_em
  7from torch_em.loss import DiceLoss
  8
  9
 10class DefaultSelfTrainingLoss(nn.Module):
 11    """Loss function for self training.
 12
 13    This loss takes as input a model and its input, as well as (pseudo) labels and potentially
 14    a mask for the labels. It then runs prediction with the model and compares the outputs
 15    to the (pseudo) labels using an internal loss function. Typically, the labels are derived
 16    from the predictions of a teacher model, and the model passed is the student model.
 17
 18    Args:
 19        loss: The internal loss function to use for comparing predictions of the teacher and student model.
 20        activation: The activation function to be applied to the prediction before passing it to the loss.
 21    """
 22    def __init__(self, loss: nn.Module = torch_em.loss.DiceLoss(), activation: Optional[nn.Module] = None):
 23        super().__init__()
 24        self.activation = activation
 25        self.loss = loss
 26        # TODO serialize the class names and kwargs instead
 27        self.init_kwargs = {}
 28
 29    def __call__(
 30        self, model: nn.Module, input_: torch.Tensor, labels: torch.Tensor, label_filter: Optional[torch.Tensor] = None
 31    ) -> torch.Tensor:
 32        """Compute the loss for self-training.
 33
 34        Args:
 35            model: The model.
 36            input_: The model inputs for this batch.
 37            labels: The (pseudo) labels for this batch.
 38            label_filter: A mask to exclude from the loss computation.
 39
 40        Returns:
 41            The loss value.
 42        """
 43        prediction = model(input_)
 44        if self.activation is not None:
 45            prediction = self.activation(prediction)
 46        if label_filter is None:
 47            loss = self.loss(prediction, labels)
 48        else:
 49            loss = self.loss(prediction * label_filter, labels * label_filter)
 50        return loss
 51
 52
 53class DefaultSelfTrainingLossAndMetric(nn.Module):
 54    """Loss and metric function for self training.
 55
 56    Similar to `DefaultSelfTrainingLoss`, but computes loss and metric value in one call
 57    to avoid running prediction with the model twice.
 58
 59    Args:
 60        loss: The internal loss function to use for comparing predictions of the teacher and student model.
 61        metric: The internal metric function to use for comparing predictions of the teacher and student model.
 62        activation: The activation function to be applied to the prediction before passing it to the loss.
 63    """
 64    def __init__(
 65        self,
 66        loss: nn.Module = torch_em.loss.DiceLoss(),
 67        metric: nn.Module = torch_em.loss.DiceLoss(),
 68        activation: Optional[nn.Module] = None
 69    ):
 70        super().__init__()
 71        self.activation = activation
 72        self.loss = loss
 73        self.metric = metric
 74        # TODO serialize the class names and dicts instead
 75        self.init_kwargs = {}
 76
 77    def __call__(self, model, input_, labels, label_filter=None):
 78        prediction = model(input_)
 79        if self.activation is not None:
 80            prediction = self.activation(prediction)
 81        if label_filter is None:
 82            loss = self.loss(prediction, labels)
 83        else:
 84            loss = self.loss(prediction * label_filter, labels * label_filter)
 85        metric = self.metric(prediction, labels)
 86        return loss, metric
 87
 88
 89def _l2_regularisation(m):
 90    l2_reg = None
 91    for W in m.parameters():
 92        if l2_reg is None:
 93            l2_reg = W.norm(2)
 94        else:
 95            l2_reg = l2_reg + W.norm(2)
 96    return l2_reg
 97
 98
 99class ProbabilisticUNetLoss(nn.Module):
100    """Training loss for ProbabilisticUNet.
101
102    Computes the ELBO loss: reconstruction term plus beta-weighted KL divergence,
103    with L2 regularisation on the posterior, prior, and fcomb weights.
104    Labels have shape (B, R, H, W), with one binary mask or class-index map per rater.
105    The ELBO averages the contributions from all raters.
106
107    Args:
108        loss: Reserved. Must be None; the ELBO objective is always used.
109    """
110    def __init__(self, loss: Optional[nn.Module] = None) -> None:
111        super().__init__()
112        self.loss = loss
113
114    def __call__(
115        self, model: nn.Module, input_: torch.Tensor, labels: torch.Tensor, label_filter: Optional[torch.Tensor] = None
116    ) -> torch.Tensor:
117        model(input_, labels)
118
119        if self.loss is None:
120            elbo = model.elbo(labels, label_filter)
121            reg_loss = (
122                _l2_regularisation(model.posterior)
123                + _l2_regularisation(model.prior)
124                + _l2_regularisation(model.fcomb.layers)
125            )
126            loss = -elbo + 1e-5 * reg_loss
127        else:
128            raise NotImplementedError("Custom loss is not supported; pass loss=None to use the ELBO.")
129
130        return loss
131
132
133class ProbabilisticUNetLossAndMetric(nn.Module):
134    """Training loss and validation metric for ProbabilisticUNet.
135
136    Computes the ELBO loss and a sample-averaged Dice metric in a single forward pass.
137    Draws prior_samples segmentation hypotheses, averages them, and evaluates against labels.
138    Labels have shape (B, R, H, W). The loss and metric average the contributions from all raters.
139
140    Args:
141        loss: Reserved. Must be None; the ELBO objective is always used.
142        metric: Metric function applied to averaged prior samples vs. labels.
143        activation: Activation applied to prior samples. The default uses sigmoid for binary outputs
144            and softmax for multiclass outputs. Pass None to use logits.
145        prior_samples: Number of prior samples to average for the metric.
146    """
147    def __init__(
148        self,
149        loss: Optional[nn.Module] = None,
150        metric: nn.Module = DiceLoss(),
151        activation: Optional[Union[nn.Module, str]] = "auto",
152        prior_samples: int = 16,
153    ) -> None:
154        super().__init__()
155        self.activation = activation
156        self.metric = metric
157        self.loss = loss
158        self.prior_samples = prior_samples
159
160    def __call__(
161        self, model: nn.Module, input_: torch.Tensor, labels: torch.Tensor, label_filter: Optional[torch.Tensor] = None
162    ):
163        model(input_, labels)
164
165        if self.loss is None:
166            elbo = model.elbo(labels, label_filter)
167            reg_loss = (
168                _l2_regularisation(model.posterior)
169                + _l2_regularisation(model.prior)
170                + _l2_regularisation(model.fcomb.layers)
171            )
172            loss = -elbo + 1e-5 * reg_loss
173        else:
174            raise NotImplementedError("Custom loss is not supported; pass loss=None to use the ELBO.")
175
176        samples_per_distribution = []
177        for _ in range(self.prior_samples):
178            samples = model.sample()
179            if self.activation == "auto":
180                samples = samples.sigmoid() if model.output_channels == 1 else samples.softmax(dim=1)
181            elif self.activation is not None:
182                samples = self.activation(samples)
183            samples_per_distribution.append(samples)
184
185        avg_samples = torch.stack(samples_per_distribution, dim=0).mean(dim=0)
186        metrics = []
187        for target in labels.split(1, dim=1):
188            if model.output_channels > 1 and isinstance(self.metric, DiceLoss):
189                target = torch.nn.functional.one_hot(target.squeeze(1).long(), model.output_channels).movedim(-1, 1)
190                target = target.float()
191            metrics.append(self.metric(avg_samples, target))
192        metric = torch.stack(metrics).mean(dim=0)
193
194        return loss, metric
195
196
197class SelfTrainingLossWithInvertibleAugmentations(nn.Module):
198    """Loss function for self-training with invertible augmentations.
199
200    Variant of `DefaultSelfTrainingLoss` for use with `FixMatchTrainerWithInvertibleAugmentations`
201    and `MeanTeacherTrainerWithInvertibleAugmentations`. Unlike `DefaultSelfTrainingLoss`, this loss
202    receives pre-computed predictions directly rather than a model and input, because the trainer
203    already applies the model and invertible augmentations before calling the loss.
204
205    Args:
206        loss: The internal loss function used to compare student predictions to pseudo-labels.
207        activation: Optional activation applied to the prediction before the loss.
208    """
209    def __init__(self, loss: nn.Module = torch_em.loss.DiceLoss(), activation: Optional[nn.Module] = None):
210        super().__init__()
211        self.activation = activation
212        self.loss = loss
213        # TODO serialize the class names and kwargs instead
214        self.init_kwargs = {}
215
216    def __call__(
217        self,
218        prediction: torch.Tensor,
219        labels: torch.Tensor,
220        label_filter: Optional[torch.Tensor] = None,
221    ) -> torch.Tensor:
222        """Compute the self-training loss.
223
224        Args:
225            prediction: Student model predictions, already mapped to the reference frame
226                via the inverse augmentation transform.
227            labels: The (pseudo) labels, mapped to the same reference frame.
228            label_filter: Optional mask or weight tensor. Where provided, both prediction
229                and labels are multiplied by this tensor before the loss is computed.
230
231        Returns:
232            The loss value.
233        """
234
235        if self.activation is not None:
236            prediction = self.activation(prediction)
237        if label_filter is None:
238            loss = self.loss(prediction, labels)
239        else:
240            loss = self.loss(prediction * label_filter, labels * label_filter)
241        return loss
242
243
244class SelfTrainingLossAndMetricWithInvertibleAugmentations(nn.Module):
245    """Loss and metric function for self-training with invertible augmentations.
246
247    Variant of `DefaultSelfTrainingLossAndMetric` for use with
248    `FixMatchTrainerWithInvertibleAugmentations` and `MeanTeacherTrainerWithInvertibleAugmentations`.
249    Computes both loss and metric in a single call from pre-computed predictions, avoiding a
250    second forward pass. Used during validation where the trainer already holds the predictions.
251
252    Args:
253        loss: The internal loss function used to compare student predictions to pseudo-labels.
254        metric: The internal metric function used to evaluate student predictions against pseudo-labels.
255        activation: Optional activation applied to the prediction before the loss and metric.
256    """
257    def __init__(
258        self,
259        loss: nn.Module = torch_em.loss.DiceLoss(),
260        metric: nn.Module = torch_em.loss.DiceLoss(),
261        activation: Optional[nn.Module] = None
262    ):
263        super().__init__()
264        self.activation = activation
265        self.loss = loss
266        self.metric = metric
267        # TODO serialize the class names and dicts instead
268        self.init_kwargs = {}
269
270    def __call__(
271        self,
272        prediction: torch.Tensor,
273        labels: torch.Tensor,
274        label_filter: Optional[torch.Tensor] = None,
275    ):
276        """Compute the self-training loss.
277
278        Args:
279            prediction: Student model predictions, already mapped to the reference frame
280                via the inverse augmentation transform.
281            labels: The (pseudo) labels, mapped to the same reference frame.
282            label_filter: Optional mask or weight tensor. Where provided, both prediction
283                and labels are multiplied by this tensor before the loss is computed.
284
285        Returns:
286            The loss and metric value.
287        """
288        if self.activation is not None:
289            prediction = self.activation(prediction)
290        if label_filter is None:
291            loss = self.loss(prediction, labels)
292        else:
293            loss = self.loss(prediction * label_filter, labels * label_filter)
294        metric = self.metric(prediction, labels)
295        return loss, metric
296
297
298class UniMatchv2Loss(nn.Module):
299    """Loss function for `UniMatchv2Trainer`.
300
301    Extends `SelfTrainingLossWithInvertibleAugmentations` to support the two-student-view scheme
302    of UniMatch v2. When `pred_dim=2`, `prediction` is expected to be a stacked tensor of two
303    student predictions `[pred_s1_inv, pred_s2_inv]`, and the loss is averaged over both views.
304    When `pred_dim=1`, it falls back to the standard single-prediction behaviour.
305
306    Args:
307        loss: The internal loss function used to compare student predictions to pseudo-labels.
308        activation: Optional activation applied to the predictions before the loss.
309    """
310    def __init__(self, loss: nn.Module = DiceLoss(), activation: Optional[nn.Module] = None):
311        super().__init__()
312        self.activation = activation
313        self.loss = loss
314        self.init_kwargs = {}
315
316    def __call__(
317        self,
318        prediction: torch.Tensor,
319        labels: torch.Tensor,
320        label_filter: Optional[torch.Tensor] = None,
321        pred_dim: int = 1,
322    ) -> torch.Tensor:
323        """Compute the UniMatch v2 self-training loss.
324
325        Args:
326            prediction: Student predictions mapped to the reference frame. When `pred_dim=2`,
327                a stacked tensor of shape `(2, B, C, ...)` holding the two strong-view predictions.
328                When `pred_dim=1`, a standard `(B, C, ...)` prediction tensor.
329            labels: The (pseudo) labels, mapped to the same reference frame.
330            label_filter: Optional mask or weight tensor applied to both prediction and labels
331                before the loss is computed.
332            pred_dim: Number of student views. Use `2` for the standard UniMatch v2 dual-view
333                training and `1` for single-view inference or validation.
334
335        Returns:
336            The loss value.
337        """
338
339        assert pred_dim in (1, 2), "pred_dim must be either 1 or 2"
340
341        if self.activation is not None:
342            prediction = self.activation(prediction)
343
344        if pred_dim == 2:
345            if label_filter is None:
346                loss = (self.loss(prediction[0], labels) + self.loss(prediction[1], labels)) / 2
347            else:
348                loss = (self.loss(
349                    prediction[0] * label_filter, labels * label_filter
350                ) + self.loss(prediction[1] * label_filter, labels * label_filter)) / 2
351            return loss
352
353        else:
354            if label_filter is None:
355                loss = self.loss(prediction, labels)
356            else:
357                loss = self.loss(prediction * label_filter, labels * label_filter)
358            return loss
359
360
361class UniMatchv2LossAndMetric(nn.Module):
362    """Loss and metric function for `UniMatchv2Trainer`.
363
364    Extends `SelfTrainingLossAndMetricWithInvertibleAugmentations` to support the two-student-view scheme
365    of UniMatch v2. `pred_dim` depends on how many views the student model processes
366    at the same time. Supports the same dual-view `pred_dim=2` convention: when two student
367    predictions are stacked, loss and metric are each averaged over both views.
368    When `pred_dim=1`, it falls back to the standard single-prediction behaviour.
369
370    Args:
371        loss: The internal loss function used to compare student predictions to pseudo-labels.
372        metric: The internal metric function used to evaluate student predictions against pseudo-labels.
373        activation: Optional activation applied to the predictions before the loss and metric.
374    """
375    def __init__(
376        self,
377        loss: nn.Module = DiceLoss(),
378        metric: nn.Module = DiceLoss(),
379        activation: Optional[nn.Module] = None
380    ):
381        super().__init__()
382        self.activation = activation
383        self.loss = loss
384        self.metric = metric
385        self.init_kwargs = {}
386
387    def __call__(
388        self,
389        prediction: torch.Tensor,
390        labels: torch.Tensor,
391        label_filter: Optional[torch.Tensor] = None,
392        pred_dim: int = 1,
393    ):
394        """Compute the UniMatch v2 self-training loss.
395
396        Args:
397            prediction: Student predictions mapped to the reference frame. When `pred_dim=2`,
398                a stacked tensor of shape `(2, B, C, ...)` holding the two strong-view predictions.
399                When `pred_dim=1`, a standard `(B, C, ...)` prediction tensor.
400            labels: The (pseudo) labels, mapped to the same reference frame.
401            label_filter: Optional mask or weight tensor applied to both prediction and labels
402                before the loss is computed.
403            pred_dim: Number of student views. Use `2` for the standard UniMatch v2 dual-view
404                training and `1` for single-view inference or validation.
405
406        Returns:
407            The loss and metric value.
408        """
409
410        if self.activation is not None:
411            prediction = self.activation(prediction)
412
413        assert pred_dim in (1, 2), "pred_dim must be either 1 or 2"
414
415        if pred_dim == 2:
416            assert len(prediction) == 2, "only implemented for list of len 2"
417            if label_filter is None:
418                loss = (self.loss(prediction[0], labels) + self.loss(prediction[1], labels)) / 2
419            else:
420                loss = (self.loss(
421                    prediction[0] * label_filter, labels * label_filter
422                ) + self.loss(prediction[1] * label_filter, labels * label_filter)) / 2
423            metric = (self.metric(prediction[0], labels) + self.metric(prediction[1], labels)) / 2
424            return loss, metric
425
426        else:
427            if label_filter is None:
428                loss = self.loss(prediction, labels)
429            else:
430                loss = self.loss(prediction * label_filter, labels * label_filter)
431            metric = self.metric(prediction, labels)
432            return loss, metric