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