torch_em.self_training.pseudo_labeling

  1from typing import Callable, Literal, Optional, Union
  2
  3import torch
  4import numpy as np
  5
  6
  7class DefaultPseudoLabeler:
  8    """Compute pseudo labels based on model predictions, typically from a teacher model.
  9
 10    Args:
 11        activation: Activation function applied to the teacher prediction.
 12        confidence_threshold: Threshold for computing a mask for filtering the pseudo labels.
 13            If None is given no mask will be computed.
 14        threshold_from_both_sides: Whether to include both values bigger than the threshold
 15            and smaller than 1 - the thrhesold, or only values bigger than the threshold, in the mask.
 16            The former should be used for binary labels, the latter for for multiclass labels.
 17        mask_channel: A specific channel to use for computing the confidence mask.
 18            By default the confidence mask is computed across all channels independently.
 19            This is useful, if only one of the channels encodes a probability.
 20    """
 21    def __init__(
 22        self,
 23        activation: Optional[torch.nn.Module] = None,
 24        confidence_threshold: Optional[float] = None,
 25        threshold_from_both_sides: bool = True,
 26        mask_channel: Optional[int] = None,
 27    ):
 28        self.activation = activation
 29        self.confidence_threshold = confidence_threshold
 30        self.threshold_from_both_sides = threshold_from_both_sides
 31        self.mask_channel = mask_channel
 32        # TODO serialize the class names and kwargs for activation instead
 33        self.init_kwargs = {
 34            "activation": None, "confidence_threshold": confidence_threshold,
 35            "threshold_from_both_sides": threshold_from_both_sides,
 36            "mask_channel": mask_channel,
 37        }
 38
 39    def _compute_label_mask_both_sides(self, pseudo_labels):
 40        upper_threshold = self.confidence_threshold
 41        lower_threshold = 1.0 - self.confidence_threshold
 42        mask = ((pseudo_labels >= upper_threshold) + (pseudo_labels <= lower_threshold)).to(dtype=torch.float32)
 43        return mask
 44
 45    def _compute_label_mask_one_side(self, pseudo_labels):
 46        mask = (pseudo_labels >= self.confidence_threshold)
 47        return mask
 48
 49    def __call__(self, teacher: torch.nn.Module, input_: torch.Tensor) -> torch.Tensor:
 50        """Compute pseudo-labels.
 51
 52        Args:
 53            teacher: The teacher model.
 54            input_: The input for this batch.
 55
 56        Returns:
 57            The pseudo-labels.
 58        """
 59        pseudo_labels = teacher(input_)
 60        if self.activation is not None:
 61            pseudo_labels = self.activation(pseudo_labels)
 62        if self.confidence_threshold is None:
 63            label_mask = None
 64        else:
 65            mask_input = pseudo_labels if self.mask_channel is None \
 66                else pseudo_labels[:, [self.mask_channel]]
 67            label_mask = self._compute_label_mask_both_sides(mask_input) if self.threshold_from_both_sides \
 68                else self._compute_label_mask_one_side(mask_input)
 69            if self.mask_channel is not None:
 70                label_mask = label_mask.expand_as(pseudo_labels)
 71        return pseudo_labels, label_mask
 72
 73    def step(self, metric, epoch):
 74        pass
 75
 76
 77class ProbabilisticPseudoLabeler:
 78    """Compute pseudo labels from a Probabilistic UNet.
 79
 80    Args:
 81        activation: Activation function applied to the teacher prediction.
 82        confidence_threshold: Threshold for computing a mask for filterign the pseudo labels.
 83            If none is given no mask will be computed.
 84        threshold_from_both_sides: Whether to include both values bigger than the threshold
 85            and smaller than 1 - the thrhesold, or only values bigger than the threshold, in the mask.
 86            The former should be used for binary labels, the latter for for multiclass labels.
 87        prior_samples: The number of times to sample from the model distribution per input.
 88        consensus_masking: Whether to activate consensus masking in the label filter.
 89            If False, the weighted consensus response (weighted per-pixel response) is returned.
 90            If True, the masked consensus response (complete aggrement of pixels) is returned.
 91    """
 92    def __init__(
 93        self,
 94        activation: Optional[torch.nn.Module] = None,
 95        confidence_threshold: Optional[float] = None,
 96        threshold_from_both_sides: bool = True,
 97        prior_samples: int = 16,
 98        consensus_masking: bool = False,
 99    ):
100        self.activation = activation
101        self.confidence_threshold = confidence_threshold
102        self.threshold_from_both_sides = threshold_from_both_sides
103        self.prior_samples = prior_samples
104        self.consensus_masking = consensus_masking
105        # TODO serialize the class names and kwargs for activation instead
106        self.init_kwargs = {
107            "activation": None, "confidence_threshold": confidence_threshold,
108            "threshold_from_both_sides": threshold_from_both_sides
109        }
110
111    def _compute_label_mask_both_sides(self, pseudo_labels):
112        upper_threshold = self.confidence_threshold
113        lower_threshold = 1.0 - self.confidence_threshold
114        mask = [
115            torch.where((sample >= upper_threshold) + (sample <= lower_threshold), torch.tensor(1.), torch.tensor(0.))
116            for sample in pseudo_labels
117        ]
118        return mask
119
120    def _compute_label_mask_one_side(self, pseudo_labels):
121        mask = [
122            torch.where((sample >= self.confidence_threshold), torch.tensor(1.), torch.tensor(0.))
123            for sample in pseudo_labels
124        ]
125        return mask
126
127    def __call__(self, teacher: torch.nn.Module, input_: torch.Tensor) -> torch.Tensor:
128        """Compute pseudo-labels.
129
130        Args:
131            teacher: The teacher model. Must be a `torch_em.model.probabilistic_unet.ProbabilisticUNet`.
132            input_: The input for this batch.
133
134        Returns:
135            The pseudo-labels.
136        """
137        teacher.forward(input_)
138        if self.activation is not None:
139            pseudo_labels = [self.activation(teacher.sample()) for _ in range(self.prior_samples)]
140        else:
141            pseudo_labels = [teacher.sample() for _ in range(self.prior_samples)]
142        pseudo_labels = torch.stack(pseudo_labels, dim=0).sum(dim=0)/self.prior_samples
143
144        if self.confidence_threshold is None:
145            label_mask = None
146        else:
147            label_mask = self._compute_label_mask_both_sides(pseudo_labels) if self.threshold_from_both_sides \
148                else self._compute_label_mask_one_side(pseudo_labels)
149            label_mask = torch.stack(label_mask, dim=0).sum(dim=0)/self.prior_samples
150            if self.consensus_masking:
151                label_mask = torch.where(label_mask == 1, 1, 0)
152
153        return pseudo_labels, label_mask
154
155    def step(self, metric, epoch):
156        pass
157
158
159class ScheduledPseudoLabeler:
160    """
161    Implement scheduled pseudo-labeling with dynamic confidence-threshold updates.
162
163    Pseudo labels are generated from a teacher model prediction and can be filtered
164    by a confidence mask. The confidence threshold can be adapted over time either
165    by decreasing it based on a monitored metric (plateau behavior) or by increasing
166    it with a fixed epoch schedule.
167
168    Args:
169        activation: Activation function applied to the teacher prediction.
170        confidence_threshold: Threshold for computing a mask for filtering the pseudo labels.
171            If None is given, no mask will be computed.
172        increase: If True, increase the confidence threshold over time according to
173            a fixed schedule. If False, decrease it based on plateau detection.
174        last_step_epoch: Last epoch at which threshold increase is allowed when
175            `increase=True`.
176        threshold_from_both_sides: Whether to include values larger than the
177            threshold and smaller than 1 - the threshold in the mask, or only values
178            larger than the threshold. The former should be used for binary labels,
179            the latter for multiclass labels.
180        mode: Determines whether the confidence threshold reduction is triggered by a "min" or "max" metric.
181            - 'min': A lower value of the monitored metric is considered better (e.g., loss).
182            - 'max': A higher value of the monitored metric is considered better (e.g., accuracy).
183        factor: Update size for confidence-threshold scheduling. Interpreted as a
184            multiplicative factor for `threshold_mode='rel'` and as an additive step
185            for `threshold_mode='abs'`.
186        patience: Number of epochs (with no improvement) after which the confidence threshold will be reduced.
187        threshold: Threshold value for determining a significant improvement in the performance metric
188            to reset the patience counter. This can be relative (percentage improvement)
189            or absolute depending on `threshold_mode`.
190        threshold_mode: Determines whether the `threshold` is interpreted as a relative improvement ('rel')
191            or an absolute improvement ('abs').
192        min_ct: Minimum allowed confidence threshold. The threshold will not be reduced below this value.
193        max_ct: Maximum allowed confidence threshold. The threshold will not be increased above this value.
194        eps: A small value to avoid floating-point precision errors during threshold comparison.
195        warm_up_epochs: Number of warm-up epochs. At the end of warm-up,
196            `confidence_threshold` is set to `max_ct`. This is intended for
197            decreasing-threshold scheduling (`increase=False`).
198        mask_channel: Specific channel to use for confidence masking. If set, the threshold
199            is applied only to that channel and the resulting mask is broadcast to all channels.
200        verbose: If True, prints messages when the confidence threshold is updated.
201    """
202
203    def __init__(
204        self,
205        activation: Optional[Union[torch.nn.Module, Callable]] = None,
206        confidence_threshold: Optional[float] = None,
207        increase: bool = False,
208        last_step_epoch: int = None,
209        threshold_from_both_sides: bool = True,
210        mode: Literal["min", "max"] = "min",
211        factor: float = 0.05,
212        patience: int = 10,
213        threshold: float = 1e-4,
214        threshold_mode: Literal["rel", "abs"] = "abs",
215        min_ct: float = 0.5,
216        max_ct: float = 0.95,
217        eps: float = 1e-8,
218        warm_up_epochs: int = 0,
219        mask_channel: Optional[int] = None,
220        verbose: bool = True,
221    ):
222        self.activation = activation
223        self.confidence_threshold = confidence_threshold
224        self.increase = increase
225        self.threshold_from_both_sides = threshold_from_both_sides
226        self.init_kwargs = {
227            "activation": None, "confidence_threshold": confidence_threshold,
228            "threshold_from_both_sides": threshold_from_both_sides,
229            "mask_channel": mask_channel,
230        }
231        # scheduler arguments
232        if mode not in {"min", "max"}:
233            raise ValueError(f"Invalid mode: {mode}. Mode should be 'min' or 'max'.")
234        self.mode = mode
235
236        assert factor < 1, f"Factor must be smaller than 1, got {factor}"
237        self.factor = factor
238
239        self.patience = patience
240        self.threshold = threshold
241
242        if threshold_mode not in {"rel", "abs"}:
243            raise ValueError(f"Invalid threshold mode: {mode}. Threshold mode should be 'rel' or 'abs'.")
244        self.threshold_mode = threshold_mode
245
246        self.min_ct = min_ct
247        self.max_ct = max_ct
248        self.eps = eps
249        self.warm_up_epochs = warm_up_epochs
250
251        if self.increase and self.warm_up_epochs > 0:
252            raise ValueError("warm_up_epochs > 0 is only supported when increase=False.")
253
254        self.mask_channel = mask_channel
255        self.verbose = verbose
256
257        if mode == "min":
258            self.best = float("inf")
259        else:  # mode == "max":
260            self.best = float("-inf")
261
262        # self.best = 0
263        self.num_bad_epochs: int = 0
264        self.last_epoch = 0
265
266        if self.increase:
267            self.last_step_epoch = last_step_epoch
268
269            n_ct = len(np.arange(self.min_ct, self.max_ct + self.factor / 2, self.factor))
270            n_increments = n_ct - 1
271
272            if n_increments <= 0:
273                # nothing to increase
274                self.step_epoch = 0
275            else:
276                # compute initial step size; enforce minimum step size of 1
277                self.step_epoch = max(1, int(np.floor(self.last_step_epoch / n_increments)))
278
279                # ensure max_ct is reachable
280                required_epochs = self.step_epoch * n_increments
281                if self.last_step_epoch < required_epochs:
282                    self.last_step_epoch = required_epochs
283
284            print(
285                f"ScheduledPseudoLabeler: Increasing confidence_threshold every {self.step_epoch} epochs;",
286                f"until epoch {self.last_step_epoch}"
287            )
288
289    def _compute_label_mask_both_sides(self, pseudo_labels):
290        upper_threshold = self.confidence_threshold
291        lower_threshold = 1.0 - self.confidence_threshold
292        mask = ((pseudo_labels >= upper_threshold) + (pseudo_labels <= lower_threshold)).to(dtype=torch.float32)
293        return mask
294
295    def _compute_label_mask_one_side(self, pseudo_labels):
296        mask = (pseudo_labels >= self.confidence_threshold)
297        return mask
298
299    def __call__(self, teacher, input_):
300        pseudo_labels = teacher(input_)
301        if self.activation is not None:
302            pseudo_labels = self.activation(pseudo_labels)
303        if self.confidence_threshold is None:
304            label_mask = None
305        else:
306            mask_input = pseudo_labels if self.mask_channel is None \
307                else pseudo_labels[:, [self.mask_channel]]
308            label_mask = self._compute_label_mask_both_sides(mask_input) if self.threshold_from_both_sides \
309                else self._compute_label_mask_one_side(mask_input)
310            if self.mask_channel is not None:
311                label_mask = label_mask.expand_as(pseudo_labels)
312        return pseudo_labels, label_mask
313
314    def end_warm_up(self):
315        self.confidence_threshold = self.max_ct
316        print(f"End of warm-up phase reached: setting confidence threshold to {self.confidence_threshold}")
317
318    def _is_better(self, a, best):
319        if self.mode == "min" and self.threshold_mode == "rel":
320            rel_epsilon = 1.0 - self.threshold
321            return a < best * rel_epsilon
322
323        elif self.mode == "min" and self.threshold_mode == "abs":
324            return a < best - self.threshold
325
326        elif self.mode == "max" and self.threshold_mode == "rel":
327            rel_epsilon = self.threshold + 1.0
328            return a > best * rel_epsilon
329
330        else:  # mode == 'max' and epsilon_mode == 'abs':
331            return a > best + self.threshold
332
333    def _reduce_ct(self, epoch):
334        old_ct = self.confidence_threshold
335        if self.threshold_mode == "rel":
336            new_ct = max(self.confidence_threshold * (1-self.factor), self.min_ct)
337        else:  # threshold_mode == 'abs':
338            new_ct = max(self.confidence_threshold - self.factor, self.min_ct)
339        if abs(old_ct - new_ct) > self.eps:
340            self.confidence_threshold = new_ct
341        if self.verbose:
342            print(f"Epoch {epoch}: reducing confidence threshold from {old_ct} to {self.confidence_threshold}")
343
344    def decrease_step(self, metric, epoch=None):
345        if epoch is None:
346            epoch = self.last_epoch + 1
347            self.last_epoch = epoch
348
349        # If the metric is None, reduce the confidence threshold every epoch
350        if metric is None:
351            if epoch == 0:
352                return
353            if epoch % self.patience == 0:
354                self._reduce_ct(epoch)
355            return
356
357        else:
358            current = float(metric)
359
360            if self._is_better(current, self.best):
361                self.best = current
362                self.num_bad_epochs = 0
363            else:
364                self.num_bad_epochs += 1
365
366            if self.num_bad_epochs > self.patience:
367                self._reduce_ct(epoch)
368                self.num_bad_epochs = 0
369
370    def _increase_ct(self, epoch):
371        old_ct = self.confidence_threshold
372        if self.threshold_mode == "rel":
373            new_ct = min(self.confidence_threshold * (1+self.factor), self.max_ct)
374        else:  # threshold_mode == 'abs':
375            new_ct = min(self.confidence_threshold + self.factor, self.max_ct)
376        if abs(old_ct - new_ct) > self.eps:
377            self.confidence_threshold = new_ct
378        if self.verbose:
379            print(f"Epoch {epoch}: increase confidence threshold from {old_ct} to {self.confidence_threshold}")
380
381    def increase_step(self, epoch):
382        if epoch > self.last_step_epoch:
383            return
384
385        if epoch % self.step_epoch == 0 and epoch != 0:
386            self._increase_ct(epoch)
387
388    def step(self, metric=None, epoch=None):
389        if epoch == self.warm_up_epochs and self.warm_up_epochs > 0:
390            self.end_warm_up()
391        elif epoch > self.warm_up_epochs:
392            if self.increase:
393                self.increase_step(epoch)
394            else:
395                self.decrease_step(metric, epoch)
396        else:
397            return