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