torch_em.self_training.fix_match

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