torch_em.self_training.mean_teacher

  1import time
  2from copy import deepcopy
  3from typing import Callable, Optional
  4
  5import torch
  6import torch_em
  7from torch_em.util import get_constructor_arguments
  8
  9from .logger import SelfTrainingTensorboardLogger
 10from ..transform.invertible_augmentations import InvertibleAugmenter
 11
 12
 13class Dummy(torch.nn.Module):
 14    init_kwargs = {}
 15
 16
 17class MeanTeacherTrainer(torch_em.trainer.DefaultTrainer):
 18    """Trainer for semi-supervised learning and domain adaptation following the MeanTeacher approach.
 19
 20    Mean Teacher was introduced by Tarvainen & Vapola in https://arxiv.org/abs/1703.01780.
 21    It uses a teacher model derived from the student model via EMA of weights
 22    to predict pseudo-labels on unlabeled data. We support two training strategies:
 23    - Joint training on labeled and unlabeled data (with a supervised and unsupervised loss function).
 24    - Training only on the unsupervised data.
 25
 26    This class expects the following data loaders:
 27    - unsupervised_train_loader: Returns two augmentations of the same input.
 28    - supervised_train_loader (optional): Returns input and labels.
 29    - unsupervised_val_loader (optional): Same as unsupervised_train_loader
 30    - supervised_val_loader (optional): Same as supervised_train_loader
 31    At least one of unsupervised_val_loader and supervised_val_loader must be given.
 32
 33    And the following elements to customize the pseudo labeling:
 34    - pseudo_labeler: to compute the psuedo-labels
 35        - Parameters: teacher, teacher_input
 36        - Returns: pseudo_labels, label_filter (<- label filter can for example be mask, weight or None)
 37    - unsupervised_loss: the loss between model predictions and pseudo labels
 38        - Parameters: model, model_input, pseudo_labels, label_filter
 39        - Returns: loss
 40    - supervised_loss (optional): the supervised loss function
 41        - Parameters: model, input, labels
 42        - Returns: loss
 43    - unsupervised_loss_and_metric (optional): the unsupervised loss function and metric
 44        - Parameters: model, model_input, pseudo_labels, label_filter
 45        - Returns: loss, metric
 46    - supervised_loss_and_metric (optional): the supervised loss function and metric
 47        - Parameters: model, input, labels
 48        - Returns: loss, metric
 49    At least one of unsupervised_loss_and_metric and supervised_loss_and_metric must be given.
 50
 51    If the parameter reinit_teacher is set to true, the teacher weights are re-initialized.
 52    If it is None, the most appropriate initialization scheme for the training approach is chosen:
 53    - semi-supervised training -> reinit, because we usually train a model from scratch
 54    - unsupervised training -> do not reinit, because we usually fine-tune a model
 55
 56    Note: adjust the batch size ratio between the 'unsupervised_train_loader' and 'supervised_train_loader'
 57    for setting the ratio between supervised and unsupervised training samples
 58
 59    Args:
 60        model: The model to be trained.
 61        unsupervised_train_loader: The loader for unsupervised training.
 62        unsupervised_loss: The loss for unsupervised training.
 63        pseudo_labeler: The pseudo labeler that predicts labels in unsupervised training.
 64        supervised_train_loader: The loader for supervised training.
 65        supervised_loss: The loss for supervised training.
 66        unsupervised_loss_and_metric: The loss and metric for unsupervised training.
 67        supervised_loss_and_metric: The loss and metrhic for supervised training.
 68        logger: The logger.
 69        momentum: The momentum value for the exponential moving weight average of the teacher model.
 70        reinit_teacher: Whether to reinit the teacher model before starting the training.
 71        sampler: A sampler for rejecting pseudo-labels according to a defined criterion.
 72        kwargs: Additional keyword arguments for `torch_em.trainer.DefaultTrainer`.
 73    """
 74
 75    def __init__(
 76        self,
 77        model: torch.nn.Module,
 78        unsupervised_train_loader: torch.utils.data.DataLoader,
 79        unsupervised_loss: torch.utils.data.DataLoader,
 80        pseudo_labeler: Callable,
 81        supervised_train_loader: Optional[torch.utils.data.DataLoader] = None,
 82        unsupervised_val_loader: Optional[torch.utils.data.DataLoader] = None,
 83        supervised_val_loader: Optional[torch.utils.data.DataLoader] = None,
 84        supervised_loss: Optional[Callable] = None,
 85        unsupervised_loss_and_metric: Optional[Callable] = None,
 86        supervised_loss_and_metric: Optional[Callable] = None,
 87        logger=SelfTrainingTensorboardLogger,
 88        momentum: float = 0.999,
 89        reinit_teacher: Optional[bool] = None,
 90        sampler: Optional[Callable] = None,
 91        **kwargs,
 92    ):
 93        self.sampler = sampler
 94        # Do we have supervised data or not?
 95        if supervised_train_loader is None:
 96            # No. -> We use the unsupervised training logic.
 97            train_loader = unsupervised_train_loader
 98            self._train_epoch_impl = self._train_epoch_unsupervised
 99        else:
100            # Yes. -> We use the semi-supervised training logic.
101            assert supervised_loss is not None
102            train_loader = supervised_train_loader if len(supervised_train_loader) < len(unsupervised_train_loader)\
103                else unsupervised_train_loader
104            self._train_epoch_impl = self._train_epoch_semisupervised
105
106        self.unsupervised_train_loader = unsupervised_train_loader
107        self.supervised_train_loader = supervised_train_loader
108
109        # Check that we have at least one of supvervised / unsupervised val loader.
110        assert sum((
111            supervised_val_loader is not None,
112            unsupervised_val_loader is not None,
113        )) > 0
114        self.supervised_val_loader = supervised_val_loader
115        self.unsupervised_val_loader = unsupervised_val_loader
116
117        if self.unsupervised_val_loader is None:
118            val_loader = self.supervised_val_loader
119        else:
120            val_loader = self.unsupervised_train_loader
121
122        # Check that we have at least one of supvervised / unsupervised loss and metric.
123        assert sum((
124            supervised_loss_and_metric is not None,
125            unsupervised_loss_and_metric is not None,
126        )) > 0
127        self.supervised_loss_and_metric = supervised_loss_and_metric
128        self.unsupervised_loss_and_metric = unsupervised_loss_and_metric
129
130        # train_loader, val_loader, loss and metric may be unnecessarily deserialized
131        kwargs.pop("train_loader", None)
132        kwargs.pop("val_loader", None)
133        kwargs.pop("metric", None)
134        kwargs.pop("loss", None)
135        super().__init__(
136            model=model, train_loader=train_loader, val_loader=val_loader,
137            loss=Dummy(), metric=Dummy(), logger=logger, **kwargs
138        )
139
140        self.unsupervised_loss = unsupervised_loss
141        self.supervised_loss = supervised_loss
142
143        self.pseudo_labeler = pseudo_labeler
144        self.momentum = momentum
145
146        # determine how we initialize the teacher weights (copy or reinitialization)
147        if reinit_teacher is None:
148            # semisupervised training: reinitialize
149            # unsupervised training: copy
150            self.reinit_teacher = supervised_train_loader is not None
151        else:
152            self.reinit_teacher = reinit_teacher
153
154        with torch.no_grad():
155            self.teacher = deepcopy(self.model)
156            if self.reinit_teacher:
157                for layer in self.teacher.children():
158                    if hasattr(layer, "reset_parameters"):
159                        layer.reset_parameters()
160            for param in self.teacher.parameters():
161                param.requires_grad = False
162
163        self._kwargs = kwargs
164
165    def _momentum_update(self):
166        # if we reinit the teacher we perform much faster updates (low momentum) in the first iterations
167        # to avoid a large gap between teacher and student weights, leading to inconsistent predictions
168        # if we don't reinit this is not necessary
169        if self.reinit_teacher:
170            current_momentum = min(1 - 1 / (self._iteration + 1), self.momentum)
171        else:
172            current_momentum = self.momentum
173
174        for param, param_teacher in zip(self.model.parameters(), self.teacher.parameters()):
175            param_teacher.data = param_teacher.data * current_momentum + param.data * (1. - current_momentum)
176
177    #
178    # functionality for saving checkpoints and initialization
179    #
180
181    def save_checkpoint(self, name, current_metric, best_metric, **extra_save_dict):
182        """@private
183        """
184        train_loader_kwargs = get_constructor_arguments(self.train_loader)
185        val_loader_kwargs = get_constructor_arguments(self.val_loader)
186        extra_state = {
187            "teacher_state": self.teacher.state_dict(),
188            "init": {
189                "train_loader_kwargs": train_loader_kwargs,
190                "train_dataset": self.train_loader.dataset,
191                "val_loader_kwargs": val_loader_kwargs,
192                "val_dataset": self.val_loader.dataset,
193                "loss_class": "torch_em.self_training.mean_teacher.Dummy",
194                "loss_kwargs": {},
195                "metric_class": "torch_em.self_training.mean_teacher.Dummy",
196                "metric_kwargs": {},
197            },
198        }
199        extra_state.update(**extra_save_dict)
200        super().save_checkpoint(name, current_metric, best_metric, **extra_state)
201
202    def load_checkpoint(self, checkpoint="best"):
203        """@private
204        """
205        save_dict = super().load_checkpoint(checkpoint)
206        self.teacher.load_state_dict(save_dict["teacher_state"])
207        self.teacher.to(self.device)
208        return save_dict
209
210    def _initialize(self, iterations, load_from_checkpoint, epochs=None):
211        best_metric = super()._initialize(iterations, load_from_checkpoint, epochs)
212        self.teacher.to(self.device)
213        return best_metric
214
215    #
216    # training and validation functionality
217    #
218
219    def _train_epoch_unsupervised(self, progress, forward_context, backprop):
220        self.model.train()
221
222        n_iter = 0
223        t_per_iter = time.time()
224
225        # Sample from both the supervised and unsupervised loader.
226        for xu1, xu2 in self.unsupervised_train_loader:
227            xu1, xu2 = xu1.to(self.device, non_blocking=True), xu2.to(self.device, non_blocking=True)
228
229            teacher_input, model_input = xu1, xu2
230
231            with forward_context(), torch.no_grad():
232                # Compute the pseudo labels.
233                pseudo_labels, label_filter = self.pseudo_labeler(self.teacher, teacher_input)
234
235            # If we have a sampler then check if the current batch matches the condition for inclusion in training.
236            if self.sampler is not None:
237                keep_batch = self.sampler(pseudo_labels, label_filter)
238                if not keep_batch:
239                    continue
240
241            self.optimizer.zero_grad()
242            # Perform unsupervised training
243            with forward_context():
244                loss = self.unsupervised_loss(self.model, model_input, pseudo_labels, label_filter)
245            backprop(loss)
246
247            if self.logger is not None:
248                with torch.no_grad(), forward_context():
249                    pred = self.model(model_input) if self._iteration % self.log_image_interval == 0 else None
250                self.logger.log_train_unsupervised(
251                    self._iteration, loss, xu1, xu2, pred, pseudo_labels, label_filter
252                )
253                lr = [pm["lr"] for pm in self.optimizer.param_groups][0]
254                self.logger.log_lr(self._iteration, lr)
255                if self.pseudo_labeler.confidence_threshold is not None:
256                    self.logger.log_ct(self._iteration, self.pseudo_labeler.confidence_threshold)
257
258            with torch.no_grad():
259                self._momentum_update()
260
261            self._iteration += 1
262            n_iter += 1
263            if self._iteration >= self.max_iteration:
264                break
265            progress.update(1)
266
267        t_per_iter = (time.time() - t_per_iter) / n_iter
268        return t_per_iter
269
270    def _train_epoch_semisupervised(self, progress, forward_context, backprop):
271        self.model.train()
272
273        n_iter = 0
274        t_per_iter = time.time()
275
276        # Sample from both the supervised and unsupervised loader.
277        for (xs, ys), (xu1, xu2) in zip(self.supervised_train_loader, self.unsupervised_train_loader):
278            xs, ys = xs.to(self.device, non_blocking=True), ys.to(self.device, non_blocking=True)
279            xu1, xu2 = xu1.to(self.device, non_blocking=True), xu2.to(self.device, non_blocking=True)
280
281            # Perform supervised training.
282            self.optimizer.zero_grad()
283            with forward_context():
284                # We pass the model, the input and the labels to the supervised loss function,
285                # so that how the loss is calculated stays flexible, e.g. to enable ELBO for PUNet.
286                supervised_loss = self.supervised_loss(self.model, xs, ys)
287
288            teacher_input, model_input = xu1, xu2
289
290            with forward_context(), torch.no_grad():
291                # Compute the pseudo labels.
292                pseudo_labels, label_filter = self.pseudo_labeler(self.teacher, teacher_input)
293
294            # Perform unsupervised training
295            with forward_context():
296                unsupervised_loss = self.unsupervised_loss(self.model, model_input, pseudo_labels, label_filter)
297
298            loss = (supervised_loss + unsupervised_loss) / 2
299            backprop(loss)
300
301            if self.logger is not None:
302                with torch.no_grad(), forward_context():
303                    unsup_pred = self.model(model_input) if self._iteration % self.log_image_interval == 0 else None
304                    supervised_pred = self.model(xs) if self._iteration % self.log_image_interval == 0 else None
305
306                self.logger.log_train_supervised(self._iteration, supervised_loss, xs, ys, supervised_pred)
307                self.logger.log_train_unsupervised(
308                    self._iteration, unsupervised_loss, xu1, xu2, unsup_pred, pseudo_labels, label_filter
309                )
310
311                self.logger.log_combined_loss(self._iteration, loss)
312                lr = [pm["lr"] for pm in self.optimizer.param_groups][0]
313                self.logger.log_lr(self._iteration, lr)
314                if self.pseudo_labeler.confidence_threshold is not None:
315                    self.logger.log_ct(self._iteration, self.pseudo_labeler.confidence_threshold)
316
317            with torch.no_grad():
318                self._momentum_update()
319
320            self._iteration += 1
321            n_iter += 1
322            if self._iteration >= self.max_iteration:
323                break
324            progress.update(1)
325
326        t_per_iter = (time.time() - t_per_iter) / n_iter
327        return t_per_iter
328
329    def _validate_supervised(self, forward_context):
330        metric_val = 0.0
331        loss_val = 0.0
332
333        for x, y in self.supervised_val_loader:
334            x, y = x.to(self.device, non_blocking=True), y.to(self.device, non_blocking=True)
335            with forward_context():
336                loss, metric = self.supervised_loss_and_metric(self.model, x, y)
337            loss_val += loss.item()
338            metric_val += metric.item()
339
340        metric_val /= len(self.supervised_val_loader)
341        loss_val /= len(self.supervised_val_loader)
342
343        if self.logger is not None:
344            with forward_context():
345                pred = self.model(x)
346            self.logger.log_validation_supervised(self._iteration, metric_val, loss_val, x, y, pred)
347
348        return metric_val
349
350    def _validate_unsupervised(self, forward_context):
351        metric_val = 0.0
352        loss_val = 0.0
353
354        for x1, x2 in self.unsupervised_val_loader:
355            x1, x2 = x1.to(self.device, non_blocking=True), x2.to(self.device, non_blocking=True)
356            teacher_input, model_input = x1, x2
357            with forward_context():
358                pseudo_labels, label_filter = self.pseudo_labeler(self.teacher, teacher_input)
359                loss, metric = self.unsupervised_loss_and_metric(self.model, model_input, pseudo_labels, label_filter)
360            loss_val += loss.item()
361            metric_val += metric.item()
362
363        metric_val /= len(self.unsupervised_val_loader)
364        loss_val /= len(self.unsupervised_val_loader)
365
366        if self.logger is not None:
367            with forward_context():
368                pred = self.model(model_input)
369            self.logger.log_validation_unsupervised(
370                self._iteration, metric_val, loss_val, x1, x2, pred, pseudo_labels, label_filter
371            )
372
373        self.pseudo_labeler.step(metric_val, self._epoch)  # NOTE: scheduler added in validation step
374
375        return metric_val
376
377    def _validate_impl(self, forward_context):
378        self.model.eval()
379
380        with torch.no_grad():
381
382            if self.supervised_val_loader is None:
383                supervised_metric = None
384            else:
385                supervised_metric = self._validate_supervised(forward_context)
386
387            if self.unsupervised_val_loader is None:
388                unsupervised_metric = None
389            else:
390                unsupervised_metric = self._validate_unsupervised(forward_context)
391
392        if unsupervised_metric is None:
393            metric = supervised_metric
394        elif supervised_metric is None:
395            metric = unsupervised_metric
396        else:
397            metric = (supervised_metric + unsupervised_metric) / 2
398
399        return metric
400
401
402class MeanTeacherTrainerWithInvertibleAugmentations(MeanTeacherTrainer):
403    """Trainer for semi-supervised learning and domain adaptation following the MeanTeacher approach.
404
405    Mean Teacher was introduced by Tarvainen & Vapola in https://arxiv.org/abs/1703.01780.
406    It uses a teacher model derived from the student model via EMA of weights
407    to predict pseudo-labels on unlabeled data. We support two training strategies:
408    - Joint training on labeled and unlabeled data (with a supervised and unsupervised loss function).
409    - Training only on the unsupervised data.
410
411    This class expects the following data loaders:
412    - unsupervised_train_loader: Returns two augmentations of the same input.
413    - supervised_train_loader (optional): Returns input and labels.
414    - unsupervised_val_loader (optional): Same as unsupervised_train_loader
415    - supervised_val_loader (optional): Same as supervised_train_loader
416    At least one of unsupervised_val_loader and supervised_val_loader must be given.
417
418    The augmenter defines separate invertible transforms for teacher and student inputs.
419    Teacher and student views are generated independently, and the corresponding inverse
420    transforms map predictions and pseudo-labels back into a shared reference frame.
421
422    And the following elements to customize the pseudo labeling:
423    - pseudo_labeler: to compute the psuedo-labels
424        - Parameters: teacher, teacher_input
425        - Returns: pseudo_labels, label_filter (<- label filter can for example be mask, weight or None)
426    - unsupervised_loss: the loss between model predictions and pseudo labels
427        - Parameters: model, model_input, pseudo_labels, label_filter
428        - Returns: loss
429    - supervised_loss (optional): the supervised loss function
430        - Parameters: model, input, labels
431        - Returns: loss
432    - unsupervised_loss_and_metric (optional): the unsupervised loss function and metric
433        - Parameters: model, model_input, pseudo_labels, label_filter
434        - Returns: loss, metric
435    - supervised_loss_and_metric (optional): the supervised loss function and metric
436        - Parameters: model, input, labels
437        - Returns: loss, metric
438    At least one of unsupervised_loss_and_metric and supervised_loss_and_metric must be given.
439
440    If the parameter reinit_teacher is set to true, the teacher weights are re-initialized.
441    If it is None, the most appropriate initialization scheme for the training approach is chosen:
442    - semi-supervised training -> reinit, because we usually train a model from scratch
443    - unsupervised training -> do not reinit, because we usually fine-tune a model
444
445    Note: adjust the batch size ratio between the 'unsupervised_train_loader' and 'supervised_train_loader'
446    for setting the ratio between supervised and unsupervised training samples
447
448    Args:
449        model: The model to be trained.
450        unsupervised_train_loader: The loader for unsupervised training.
451        unsupervised_loss: The loss for unsupervised training.
452        pseudo_labeler: The pseudo labeler that predicts labels in unsupervised training.
453        augmenter: Invertible augmenter providing separate teacher and student transforms,
454            including corresponding inverse transforms to align predictions in a common frame.
455        supervised_train_loader: The loader for supervised training.
456        supervised_loss: The loss for supervised training.
457        unsupervised_loss_and_metric: The loss and metric for unsupervised training.
458        supervised_loss_and_metric: The loss and metrhic for supervised training.
459        logger: The logger.
460        momentum: The momentum value for the exponential moving weight average of the teacher model.
461        reinit_teacher: Whether to reinit the teacher model before starting the training.
462        sampler: A sampler for rejecting pseudo-labels according to a defined criterion.
463        kwargs: Additional keyword arguments for `torch_em.trainer.DefaultTrainer`.
464    """
465
466    def __init__(
467        self,
468        model: torch.nn.Module,
469        unsupervised_train_loader: torch.utils.data.DataLoader,
470        unsupervised_loss: torch.utils.data.DataLoader,
471        pseudo_labeler: Callable,
472        augmenter: InvertibleAugmenter,
473        supervised_train_loader: Optional[torch.utils.data.DataLoader] = None,
474        unsupervised_val_loader: Optional[torch.utils.data.DataLoader] = None,
475        supervised_val_loader: Optional[torch.utils.data.DataLoader] = None,
476        supervised_loss: Optional[Callable] = None,
477        unsupervised_loss_and_metric: Optional[Callable] = None,
478        supervised_loss_and_metric: Optional[Callable] = None,
479        logger=SelfTrainingTensorboardLogger,
480        momentum: float = 0.999,
481        reinit_teacher: Optional[bool] = None,
482        sampler: Optional[Callable] = None,
483        **kwargs,
484    ):
485        super().__init__(
486            model=model,
487            unsupervised_train_loader=unsupervised_train_loader,
488            unsupervised_loss=unsupervised_loss,
489            pseudo_labeler=pseudo_labeler,
490            supervised_train_loader=supervised_train_loader,
491            unsupervised_val_loader=unsupervised_val_loader,
492            supervised_val_loader=supervised_val_loader,
493            supervised_loss=supervised_loss,
494            unsupervised_loss_and_metric=unsupervised_loss_and_metric,
495            supervised_loss_and_metric=supervised_loss_and_metric,
496            logger=logger,
497            momentum=momentum,
498            reinit_teacher=reinit_teacher,
499            sampler=sampler,
500            **kwargs,
501        )
502        self.augmenter = augmenter
503
504    #
505    # training and validation functionality
506    #
507
508    def _train_epoch_unsupervised(self, progress, forward_context, backprop):
509        self.model.train()
510
511        n_iter = 0
512        t_per_iter = time.time()
513
514        # Sample from both the supervised and unsupervised loader.
515        for xu in self.unsupervised_train_loader:
516            self.augmenter.reset_all()
517            xu = xu.to(self.device, non_blocking=True)
518
519            xu1, xu2 = self.augmenter.teacher.transform(xu), self.augmenter.student.transform(xu)
520            teacher_input, model_input = xu1, xu2
521
522            with forward_context(), torch.no_grad():
523                # Compute the pseudo labels.
524                pseudo_labels, label_filter = self.pseudo_labeler(self.teacher, teacher_input)
525                pseudo_labels_inv = self.augmenter.teacher.reverse_transform(pseudo_labels)
526                label_filter_inv = (
527                    self.augmenter.teacher.reverse_transform(label_filter)
528                    if label_filter is not None else None
529                )
530
531            # If we have a sampler then check if the current batch matches the condition for inclusion in training.
532            if self.sampler is not None:
533                keep_batch = self.sampler(pseudo_labels, label_filter)
534                if not keep_batch:
535                    continue
536
537            self.optimizer.zero_grad()
538            # Perform unsupervised training
539            with forward_context():
540                pred = self.model(model_input)
541                pred_inv = self.augmenter.student.reverse_transform(pred)
542                loss = self.unsupervised_loss(pred_inv, pseudo_labels_inv, label_filter_inv)
543            backprop(loss)
544
545            if self.logger is not None:
546                with torch.no_grad(), forward_context():
547                    pred = pred if self._iteration % self.log_image_interval == 0 else None
548                self.logger.log_train_unsupervised(
549                    self._iteration, loss, xu, xu, pred_inv, pseudo_labels_inv, label_filter_inv
550                )
551                self.logger.log_train_augmentations(
552                    self._iteration, xu1, xu2, pseudo_labels, pred,
553                )
554                lr = [pm["lr"] for pm in self.optimizer.param_groups][0]
555                self.logger.log_lr(self._iteration, lr)
556                if self.pseudo_labeler.confidence_threshold is not None:
557                    self.logger.log_ct(self._iteration, self.pseudo_labeler.confidence_threshold)
558
559            with torch.no_grad():
560                self._momentum_update()
561
562            self._iteration += 1
563            n_iter += 1
564            if self._iteration >= self.max_iteration:
565                break
566            progress.update(1)
567
568        t_per_iter = (time.time() - t_per_iter) / n_iter
569        return t_per_iter
570
571    def _train_epoch_semisupervised(self, progress, forward_context, backprop):
572        self.model.train()
573
574        n_iter = 0
575        t_per_iter = time.time()
576
577        # Sample from both the supervised and unsupervised loader.
578        for (xs, ys), xu in zip(self.supervised_train_loader, self.unsupervised_train_loader):
579            self.augmenter.reset_all()
580            xs, ys = xs.to(self.device, non_blocking=True), ys.to(self.device, non_blocking=True)
581            xu = xu.to(self.device, non_blocking=True)
582
583            xu1, xu2 = self.augmenter.teacher.transform(xu), self.augmenter.student.transform(xu)
584            teacher_input, model_input = xu1, xu2
585
586            # Perform supervised training.
587            self.optimizer.zero_grad()
588            with forward_context():
589                # We pass the model, the input and the labels to the supervised loss function,
590                # so that how the loss is calculated stays flexible, e.g. to enable ELBO for PUNet.
591                supervised_pred = self.model(xs)
592                supervised_loss = self.supervised_loss(supervised_pred, ys)
593
594            with forward_context(), torch.no_grad():
595                # Compute the pseudo labels.
596                pseudo_labels, label_filter = self.pseudo_labeler(self.teacher, teacher_input)
597                pseudo_labels_inv = self.augmenter.teacher.reverse_transform(pseudo_labels)
598                label_filter_inv = (
599                    self.augmenter.teacher.reverse_transform(label_filter)
600                    if label_filter is not None else None
601                )
602
603            # Perform unsupervised training
604            with forward_context():
605                unsup_pred = self.model(model_input)
606                unsup_pred_inv = self.augmenter.student.reverse_transform(unsup_pred)
607                unsupervised_loss = self.unsupervised_loss(unsup_pred_inv, pseudo_labels_inv, label_filter_inv)
608
609            loss = (supervised_loss + unsupervised_loss) / 2
610            backprop(loss)
611
612            if self.logger is not None:
613                with torch.no_grad(), forward_context():
614                    unsup_pred = unsup_pred if self._iteration % self.log_image_interval == 0 else None
615                    supervised_pred = supervised_pred if self._iteration % self.log_image_interval == 0 else None
616
617                self.logger.log_train_supervised(self._iteration, supervised_loss, xs, ys, supervised_pred)
618                self.logger.log_train_unsupervised(
619                    self._iteration, unsupervised_loss, xu, xu, unsup_pred_inv, pseudo_labels_inv, label_filter_inv
620                )
621                self.logger.log_train_augmentations(
622                    self._iteration, xu1, xu2, pseudo_labels, unsup_pred,
623                )
624
625                self.logger.log_combined_loss(self._iteration, loss)
626                lr = [pm["lr"] for pm in self.optimizer.param_groups][0]
627                self.logger.log_lr(self._iteration, lr)
628                if self.pseudo_labeler.confidence_threshold is not None:
629                    self.logger.log_ct(self._iteration, self.pseudo_labeler.confidence_threshold)
630
631            with torch.no_grad():
632                self._momentum_update()
633
634            self._iteration += 1
635            n_iter += 1
636            if self._iteration >= self.max_iteration:
637                break
638            progress.update(1)
639
640        t_per_iter = (time.time() - t_per_iter) / n_iter
641        return t_per_iter
642
643    def _validate_supervised(self, forward_context):
644        metric_val = 0.0
645        loss_val = 0.0
646
647        for x, y in self.supervised_val_loader:
648            x, y = x.to(self.device, non_blocking=True), y.to(self.device, non_blocking=True)
649            with forward_context():
650                pred = self.model(x)
651                loss, metric = self.supervised_loss_and_metric(pred, y)
652            loss_val += loss.item()
653            metric_val += metric.item()
654
655        metric_val /= len(self.supervised_val_loader)
656        loss_val /= len(self.supervised_val_loader)
657
658        if self.logger is not None:
659            self.logger.log_validation_supervised(self._iteration, metric_val, loss_val, x, y, pred)
660
661        return metric_val
662
663    def _validate_unsupervised(self, forward_context):
664        metric_val = 0.0
665        loss_val = 0.0
666
667        for x in self.unsupervised_val_loader:
668            self.augmenter.reset_all()
669            x = x.to(self.device, non_blocking=True)
670            x1, x2 = self.augmenter.teacher.transform(x), self.augmenter.student.transform(x)
671            teacher_input, model_input = x1, x2
672            with forward_context():
673                pseudo_labels, label_filter = self.pseudo_labeler(self.teacher, teacher_input)
674                pseudo_labels_inv = self.augmenter.teacher.reverse_transform(pseudo_labels)
675                label_filter_inv = (
676                    self.augmenter.teacher.reverse_transform(label_filter)
677                    if label_filter is not None else None
678                )
679
680                pred = self.model(model_input)
681                pred_inv = self.augmenter.student.reverse_transform(pred)
682                loss, metric = self.unsupervised_loss_and_metric(pred_inv, pseudo_labels_inv, label_filter_inv)
683            loss_val += loss.item()
684            metric_val += metric.item()
685
686        metric_val /= len(self.unsupervised_val_loader)
687        loss_val /= len(self.unsupervised_val_loader)
688
689        if self.logger is not None:
690            self.logger.log_validation_unsupervised(
691                self._iteration, metric_val, loss_val, x, x, pred_inv, pseudo_labels_inv, label_filter_inv
692            )
693            self.logger.log_validation_augmentations(
694                self._iteration, x1, x2, pseudo_labels, pred,
695            )
696
697        self.pseudo_labeler.step(metric_val, self._epoch)
698
699        return metric_val
700
701    def _validate_impl(self, forward_context):
702        self.model.eval()
703
704        with torch.no_grad():
705
706            if self.supervised_val_loader is None:
707                supervised_metric = None
708            else:
709                supervised_metric = self._validate_supervised(forward_context)
710
711            if self.unsupervised_val_loader is None:
712                unsupervised_metric = None
713            else:
714                unsupervised_metric = self._validate_unsupervised(forward_context)
715
716        if unsupervised_metric is None:
717            metric = supervised_metric
718        elif supervised_metric is None:
719            metric = unsupervised_metric
720        else:
721            metric = (supervised_metric + unsupervised_metric) / 2
722
723        return metric