torch_em.self_training.probabilistic_unet_trainer

  1import time
  2import torch
  3import torch_em
  4
  5
  6class DummyLoss(torch.nn.Module):
  7    pass
  8
  9
 10class ProbabilisticUNetTrainer(torch_em.trainer.DefaultTrainer):
 11    """This trainer implements training for the 'Probabilistic UNet' of Kohl et al.: (https://arxiv.org/abs/1806.05034).
 12    This approach combines the learnings from UNet and VAEs (Prior and Posterior networks) to obtain generative
 13    segmentations. The heuristic trains by taking into account the feature maps from UNet and the samples from
 14    the posterior distribution, estimating the loss and further sampling from the prior for validation.
 15
 16    Parameters:
 17        clipping_value [float] - (default: None)
 18        prior_samples [int] - (default: 16)
 19        loss [callable] - (default: None)
 20        loss_and_metric [callable] - (default: None)
 21    """
 22
 23    def __init__(
 24            self,
 25            clipping_value=None,
 26            prior_samples=16,
 27            loss=None,
 28            loss_and_metric=None,
 29            **kwargs
 30    ):
 31        super().__init__(loss=loss, metric=DummyLoss(), **kwargs)
 32        assert loss, loss_and_metric is not None
 33
 34        self.loss_and_metric = loss_and_metric
 35
 36        self.clipping_value = clipping_value
 37
 38        self.prior_samples = prior_samples
 39        self.sigmoid = torch.nn.Sigmoid()
 40
 41        self._kwargs = kwargs
 42
 43    #
 44    # functionality for sampling from the network
 45    #
 46
 47    def _sample(self):
 48        samples = [self.model.sample() for _ in range(self.prior_samples)]
 49        return samples
 50
 51    #
 52    # training and validation functionality
 53    #
 54
 55    def _train_epoch_impl(self, progress, forward_context, backprop):
 56        self.model.train()
 57
 58        n_iter = 0
 59        t_per_iter = time.time()
 60
 61        for x, y in self.train_loader:
 62            x, y = x.to(self.device), y.to(self.device)
 63
 64            self.optimizer.zero_grad()
 65
 66            with forward_context():
 67                # We pass the model, the input and the labels to the supervised loss function, so
 68                # that's how the loss is calculated stays flexible, e.g. here to enable ELBO for PUNet.
 69                loss = self.loss(self.model, x, y)
 70
 71            backprop(loss)
 72
 73            # To counter the exploding gradients in the posterior net
 74            if self.clipping_value is not None:
 75                torch.nn.utils.clip_grad_norm_(self.model.posterior.encoder.layers.parameters(), self.clipping_value)
 76
 77            if self.logger is not None:
 78                lr = [pm["lr"] for pm in self.optimizer.param_groups][0]
 79                samples = self._sample() if self._iteration % self.log_image_interval == 0 else None
 80                self.logger.log_train(self._iteration, loss, lr, x, y, samples)
 81
 82            self._iteration += 1
 83            n_iter += 1
 84            if self._iteration >= self.max_iteration:
 85                break
 86            progress.update(1)
 87
 88        t_per_iter = (time.time() - t_per_iter) / n_iter
 89        return t_per_iter
 90
 91    def _validate_impl(self, forward_context):
 92        self.model.eval()
 93
 94        metric_val = 0.0
 95        loss_val = 0.0
 96
 97        with torch.no_grad():
 98            for x, y in self.val_loader:
 99                x, y = x.to(self.device), y.to(self.device)
100
101                with forward_context():
102                    loss, metric = self.loss_and_metric(self.model, x, y)
103
104                loss_val += loss.item()
105                metric_val += metric
106
107        metric_val /= len(self.val_loader)
108        loss_val /= len(self.val_loader)
109
110        if self.logger is not None:
111            samples = self._sample()
112            self.logger.log_validation(self._iteration, metric_val, loss_val, x, y, samples)
113
114        return metric_val