torch_em.trainer.spoco_trainer

  1import time
  2from copy import deepcopy
  3
  4import torch
  5from .default_trainer import DefaultTrainer
  6from .tensorboard_logger import TensorboardLogger
  7
  8
  9class SPOCOTrainer(DefaultTrainer):
 10    def __init__(
 11        self,
 12        model,
 13        momentum=0.999,
 14        semisupervised_loss=None,
 15        semisupervised_loader=None,
 16        logger=TensorboardLogger,
 17        **kwargs,
 18    ):
 19        super().__init__(model=model, logger=logger, **kwargs)
 20        self.momentum = momentum
 21        # copy the model and don"t require gradients for it
 22        self.model2 = deepcopy(self.model)
 23        for param in self.model2.parameters():
 24            param.requires_grad = False
 25        # do we have a semi-supervised loss and loader?
 26        assert (semisupervised_loss is None) == (semisupervised_loader is None)
 27        self.semisupervised_loader = semisupervised_loader
 28        self.semisupervised_loss = semisupervised_loss
 29        self._kwargs = kwargs
 30
 31    def _momentum_update(self):
 32        for param_model, param_teacher in zip(self.model.parameters(), self.model2.parameters()):
 33            param_teacher.data = param_teacher.data * self.momentum + param_model.data * (1. - self.momentum)
 34
 35    def save_checkpoint(self, name, current_metric, best_metric, **extra_save_dict):
 36        super().save_checkpoint(
 37            name, current_metric, best_metric, model2_state=self.model2.state_dict(), **extra_save_dict
 38        )
 39
 40    def load_checkpoint(self, checkpoint="best"):
 41        save_dict = super().load_checkpoint(checkpoint)
 42        self.model2.load_state_dict(save_dict["model2_state"])
 43        self.model2.to(self.device)
 44        return save_dict
 45
 46    def _initialize(self, iterations, load_from_checkpoint, epochs=None):
 47        best_metric = super()._initialize(iterations, load_from_checkpoint, epochs)
 48        self.model2.to(self.device)
 49        return best_metric
 50
 51    def _train_epoch_semisupervised(self, progress, forward_context, backprop):
 52        self.model.train()
 53        self.model2.train()
 54        progress.set_description(
 55            f"Run semi-supervised training for {len(self.semisupervised_loader)} iterations", refresh=True
 56        )
 57
 58        for x in self.semisupervised_loader:
 59            x = x.to(self.device)
 60            self.optimizer.zero_grad()
 61
 62            with forward_context():
 63                prediction = self.model(x)
 64                with torch.no_grad():
 65                    prediction2 = self.model2(x)
 66                loss = self.semisupervised_loss(prediction, prediction2)
 67            backprop(loss)
 68
 69            with torch.no_grad():
 70                self._momentum_update()
 71
 72    def _train_epoch_impl(self, progress, forward_context, backprop):
 73        self.model.train()
 74        self.model2.train()
 75
 76        n_iter = 0
 77        t_per_iter = time.time()
 78        for x, y in self.train_loader:
 79            x, y = x.to(self.device), y.to(self.device)
 80
 81            self.optimizer.zero_grad()
 82
 83            with forward_context():
 84                prediction = self.model(x)
 85                with torch.no_grad():
 86                    prediction2 = self.model2(x)
 87                loss = self.loss((prediction, prediction2), y)
 88
 89            if self._iteration % self.log_image_interval == 0:
 90                prediction.retain_grad()
 91
 92            backprop(loss)
 93
 94            with torch.no_grad():
 95                self._momentum_update()
 96
 97            lr = [pm["lr"] for pm in self.optimizer.param_groups][0]
 98            if self.logger is not None:
 99                self.logger.log_train(self._iteration, loss, lr,
100                                      x, y, prediction,
101                                      log_gradients=True)
102
103            self._iteration += 1
104            n_iter += 1
105            if self._iteration >= self.max_iteration:
106                break
107            progress.update(1)
108
109        if self.semisupervised_loader is not None:
110            self._train_epoch_semisupervised(progress, forward_context, backprop)
111        t_per_iter = (time.time() - t_per_iter) / n_iter
112        return t_per_iter
113
114    def _validate_impl(self, forward_context):
115        self.model.eval()
116        self.model2.eval()
117
118        metric = 0.0
119        loss = 0.0
120
121        with torch.no_grad():
122            for x, y in self.val_loader:
123                x, y = x.to(self.device), y.to(self.device)
124                with forward_context():
125                    prediction = self.model(x)
126                    prediction2 = self.model2(x)
127                loss += self.loss((prediction, prediction2), y).item()
128                metric += self.metric(prediction, y).item()
129
130        metric /= len(self.val_loader)
131        loss /= len(self.val_loader)
132        if self.logger is not None:
133            self.logger.log_validation(self._iteration, metric, loss, x, y, prediction)
134        return metric