torch_em.self_training.uni_match_v2

  1import time
  2
  3import torch
  4from torch_em.self_training.mean_teacher import MeanTeacherTrainerWithInvertibleAugmentations
  5
  6# os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
  7
  8
  9class UniMatchv2Trainer(MeanTeacherTrainerWithInvertibleAugmentations):
 10    """Trainer for semi-supervised learning and domain adaptation following the UniMatch v2 framework.
 11
 12    UniMatch v2 was introduced by Yang et al. in https://arxiv.org/abs/2410.10777v2.
 13    It uses a teacher model derived from the student model via EMA of weights to predict
 14    pseudo-labels on unlabeled data. Three augmented views are generated per sample - one weak
 15    (for the teacher) and two strong (for the student) - and the student loss is computed as the
 16    average over both strong-view predictions against the shared weak-view pseudo-label.
 17    We support two training strategies:
 18    - Joint training on labeled and unlabeled data (with a supervised and unsupervised loss function).
 19    - Training only on the unsupervised data.
 20
 21    This class expects the following data loaders:
 22    - unsupervised_train_loader: Returns a single (raw) input per sample. The trainer applies
 23      weak and two strong augmentations internally via the augmenter.
 24    - supervised_train_loader (optional): Returns input and labels.
 25    - unsupervised_val_loader (optional): Same format as unsupervised_train_loader.
 26    - supervised_val_loader (optional): Same format as supervised_train_loader.
 27    At least one of unsupervised_val_loader and supervised_val_loader must be given.
 28
 29    The augmenter must be a `UniMatchv2Augmenters` instance providing three invertible transforms:
 30    `.weak` for the teacher view, `.strong1` and `.strong2` for the two student views. The
 31    corresponding inverse transforms map predictions and pseudo-labels back into a shared
 32    reference frame before the loss is computed.
 33
 34    The following arguments can be used to customize the pseudo labeling:
 35    - pseudo_labeler: to compute the pseudo-labels
 36        - Parameters: teacher, teacher_input
 37        - Returns: pseudo_labels, label_filter (<- label filter can for example be mask, weight or None)
 38    - unsupervised_loss: the loss between stacked student predictions and pseudo-labels
 39        - Parameters: prediction (stacked [pred_s1_inv, pred_s2_inv]), pseudo_labels, label_filter, pred_dim
 40        - Returns: loss
 41    - supervised_loss (optional): the supervised loss function
 42        - Parameters: prediction, labels
 43        - Returns: loss
 44    - unsupervised_loss_and_metric (optional): the unsupervised loss function and metric
 45        - Parameters: prediction (stacked), pseudo_labels, label_filter, pred_dim
 46        - Returns: loss, metric
 47    - supervised_loss_and_metric (optional): the supervised loss function and metric
 48        - Parameters: prediction, labels
 49        - Returns: loss, metric
 50    At least one of unsupervised_loss_and_metric and supervised_loss_and_metric must be given.
 51
 52    Note: adjust the batch size of the 'unsupervised_train_loader' relative to
 53    'supervised_train_loader' to control the ratio of supervised to unsupervised training samples.
 54
 55    Args:
 56        model: The model to be trained.
 57        unsupervised_train_loader: The loader for unsupervised training (returns raw inputs only).
 58        unsupervised_loss: The loss for unsupervised training.
 59        pseudo_labeler: The pseudo labeler that predicts labels in unsupervised training.
 60        augmenter: `UniMatchv2Augmenters` instance providing `.weak`, `.strong1`, and `.strong2`
 61            invertible transforms with corresponding inverse transforms.
 62        complementary_dropout: If True, applies complementary feature dropout to the encoder
 63            features before decoding, creating two complementary student views. Requires a
 64            UNETR-compatible model architecture.
 65        supervised_train_loader: The loader for supervised training.
 66        supervised_loss: The loss for supervised training.
 67        unsupervised_loss_and_metric: The loss and metric for unsupervised training.
 68        supervised_loss_and_metric: The loss and metric for supervised training.
 69        logger: The logger. Defaults to `UniMatchv2TensorboardLogger`.
 70        momentum: The momentum value for the exponential moving weight average of the teacher model.
 71        reinit_teacher: Whether to reinit the teacher model before starting the training.
 72        sampler: A sampler for rejecting pseudo-labels according to a defined criterion.
 73        separate_backward: Whether to backpropagate each loss term separately before one optimizer update.
 74            The default uses separate optimizer updates for the supervised and unsupervised losses.
 75        kwargs: Additional keyword arguments for `torch_em.trainer.DefaultTrainer`.
 76    """
 77
 78    def __init__(self, complementary_dropout, separate_backward: bool = False, **kwargs):
 79        super().__init__(separate_backward=separate_backward, **kwargs)
 80        self.complementary_dropout = complementary_dropout
 81
 82        self.teacher.eval()
 83
 84    def unetr_decoder_prediction(self, model, features, input_shape, original_shape):
 85
 86        z9 = model.deconv1(features)
 87        z6 = model.deconv2(z9)
 88        z3 = model.deconv3(z6)
 89        z0 = model.deconv4(z3)
 90
 91        updated_from_encoder = [z9, z6, z3]
 92
 93        x = model.base(features)
 94        x = model.decoder(x, encoder_inputs=updated_from_encoder)
 95        x = model.deconv_out(x)
 96
 97        x = torch.cat([x, z0], dim=1)
 98        x = model.decoder_head(x)
 99
100        x = model.out_conv(x)
101        if model.final_activation is not None:
102            x = model.final_activation(x)
103
104        x = model.postprocess_masks(x, input_shape, original_shape)
105        return x
106
107    def predict_with_comp_drop(self, model, input_):
108        batch_size = input_.shape[0]
109        original_shape = input_.shape[2:]
110
111        x, input_shape = model.preprocess(input_)
112
113        if len(original_shape) == 2:
114            encoder_output = model.encoder(x)
115            if isinstance(encoder_output[-1], list):
116                features, _ = encoder_output
117            else:
118                features = encoder_output
119        if len(original_shape) == 3:
120            depth = input_.shape[-3]
121            features = torch.stack([model.encoder(x[:, :, i])[0] for i in range(depth)], dim=2)
122
123        features_dim = features.shape[1]
124
125        binom = torch.distributions.binomial.Binomial(probs=0.5)
126
127        dropout_mask1 = binom.sample((int(batch_size/2), features_dim)).to(input_.device) * 2.0
128        if len(original_shape) == 2:
129            dropout_mask1 = dropout_mask1.unsqueeze(-1).unsqueeze(-1)
130        if len(original_shape) == 3:
131            dropout_mask1 = dropout_mask1.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
132
133        dropout_mask2 = 2.0 - dropout_mask1
134        dropout_mask = torch.cat([dropout_mask1, dropout_mask2])
135
136        # NOTE: in the UniMatch v2 code some samples of the batch stay unchanged!
137        # Keep some samples unchanged (code block not tested)
138        # dropout_prob = 0.5
139        # num_kept = int(batch_size * (1 - dropout_prob))
140        # kept_indexes = torch.randperm(batch_size, device=input_.device)[:num_kept]
141
142        # dropout_mask1[kept_indexes, :] = 1.0
143        # dropout_mask2[kept_indexes, :] = 1.0
144
145        dropped_features = features * dropout_mask
146
147        pred = self.unetr_decoder_prediction(model, dropped_features, input_shape, original_shape)
148
149        return pred
150
151    def _train_epoch_unsupervised(
152        self, progress, forward_context, backprop
153    ):
154        self.model.train()
155
156        n_iter = 0
157        t_per_iter = time.time()
158
159        for x_u in self.unsupervised_train_loader:
160            self.augmenter.reset_all()
161            x_u = x_u.to(self.device, non_blocking=True)
162
163            x_u_w = self.augmenter.weak.transform(x_u)
164            x_u_s1, x_u_s2 = self.augmenter.strong1.transform(x_u), self.augmenter.strong2.transform(x_u)
165
166            # Compute the pseudo labels (unsupervised teacher prediction)
167            with forward_context(), torch.no_grad():
168                pseudo_labels, label_filter = self.pseudo_labeler(self.teacher, x_u_w)
169                pseudo_labels_inv = self.augmenter.weak.reverse_transform(pseudo_labels)
170                label_filter_inv = (
171                    self.augmenter.weak.reverse_transform(label_filter)
172                    if label_filter is not None else None
173                )
174
175            # Perform unsupervised training
176            with forward_context():
177                if self.complementary_dropout:
178                    pred_s1, pred_s2 = self.predict_with_comp_drop(self.model, torch.cat((x_u_s1, x_u_s2))).chunk(2)
179                else:
180                    pred_s1, pred_s2 = self.model(torch.cat((x_u_s1, x_u_s2))).chunk(2)
181                pred_s1_inv = self.augmenter.strong1.reverse_transform(pred_s1)
182                pred_s2_inv = self.augmenter.strong2.reverse_transform(pred_s2)
183                unsupervised_loss = self.unsupervised_loss(
184                    torch.stack((pred_s1_inv, pred_s2_inv)),
185                    pseudo_labels_inv,
186                    label_filter_inv,
187                    pred_dim=2,
188                )
189
190            backprop(unsupervised_loss)
191
192            if self.logger is not None:
193                self.logger.log_train_unsupervised(
194                    self._iteration,
195                    unsupervised_loss,
196                    x_u,
197                    pred_s1_inv,
198                    pred_s2_inv,
199                    pseudo_labels_inv,
200                    label_filter_inv,
201                )
202                self.logger.log_train_augmentations(
203                    self._iteration,
204                    x_u_w,
205                    x_u_s1,
206                    x_u_s2,
207                    pseudo_labels,
208                    pred_s1,
209                    pred_s2,
210                )
211
212                lr = [pm["lr"] for pm in self.optimizer.param_groups][0]
213                self.logger.log_lr(self._iteration, lr)
214                if self.pseudo_labeler.confidence_threshold is not None:
215                    self.logger.log_ct(self._iteration, self.pseudo_labeler.confidence_threshold)
216
217            with torch.no_grad():
218                self._momentum_update()  # EMA update of the teacher
219
220            self._iteration += 1
221            n_iter += 1
222            if self._iteration >= self.max_iteration:
223                break
224            progress.update(1)
225
226        t_per_iter = (time.time() - t_per_iter) / n_iter
227        return t_per_iter
228
229    def _train_epoch_semisupervised(
230        self, progress, forward_context, backprop
231    ):
232        train_loader = zip(self.supervised_train_loader, self.unsupervised_train_loader)
233        self.model.train()
234
235        n_iter = 0
236        t_per_iter = time.time()
237
238        for i, ((x_s, y_s), x_u) in enumerate(train_loader):
239            self.augmenter.reset_all()
240
241            x_s, y_s = x_s.to(self.device, non_blocking=True), y_s.to(self.device, non_blocking=True)
242            x_u = x_u.to(self.device, non_blocking=True)
243
244            x_u_w = self.augmenter.weak.transform(x_u)
245            x_u_s1, x_u_s2 = self.augmenter.strong1.transform(x_u), self.augmenter.strong2.transform(x_u)
246
247            self.optimizer.zero_grad()
248            # supervised loss (supervised student prediction)
249            with forward_context():
250                pred_s = self.model(x_s)
251                supervised_loss = self.supervised_loss(pred_s, y_s)
252
253            if self.separate_backward:
254                self._accumulate_gradients(supervised_loss)
255            else:
256                backprop(supervised_loss)
257
258            # Compute the pseudo labels (unsupervised teacher prediction)
259            with forward_context(), torch.no_grad():
260                pseudo_labels, label_filter = self.pseudo_labeler(self.teacher, x_u_w)
261                pseudo_labels_inv = self.augmenter.weak.reverse_transform(pseudo_labels)
262                label_filter_inv = (
263                    self.augmenter.weak.reverse_transform(label_filter)
264                    if label_filter is not None else None
265                )
266
267            # Perform unsupervised training
268            if not self.separate_backward:
269                self.optimizer.zero_grad()
270            with forward_context():
271                if self.complementary_dropout:
272                    pred_s1, pred_s2 = self.predict_with_comp_drop(self.model, torch.cat((x_u_s1, x_u_s2))).chunk(2)
273                else:
274                    pred_s1, pred_s2 = self.model(torch.cat((x_u_s1, x_u_s2))).chunk(2)
275                pred_s1_inv = self.augmenter.strong1.reverse_transform(pred_s1)
276                pred_s2_inv = self.augmenter.strong2.reverse_transform(pred_s2)
277                unsupervised_loss = self.unsupervised_loss(
278                    torch.stack((pred_s1_inv, pred_s2_inv)),
279                    pseudo_labels_inv,
280                    label_filter_inv,
281                    pred_dim=2,
282                )
283
284            backprop(unsupervised_loss)
285
286            if self.logger is not None:
287                self.logger.log_train_supervised(
288                    self._iteration, supervised_loss, x_s, y_s, pred_s
289                )
290                self.logger.log_train_unsupervised(
291                    self._iteration,
292                    unsupervised_loss,
293                    x_u,
294                    pred_s1_inv,
295                    pred_s2_inv,
296                    pseudo_labels_inv,
297                    label_filter_inv,
298                )
299                self.logger.log_train_augmentations(
300                    self._iteration,
301                    x_u_w,
302                    x_u_s1,
303                    x_u_s2,
304                    pseudo_labels,
305                    pred_s1,
306                    pred_s2,
307                )
308
309                lr = [pm["lr"] for pm in self.optimizer.param_groups][0]
310                self.logger.log_lr(self._iteration, lr)
311                if self.pseudo_labeler.confidence_threshold is not None:
312                    self.logger.log_ct(self._iteration, self.pseudo_labeler.confidence_threshold)
313
314            with torch.no_grad():
315                self._momentum_update()  # EMA update of the teacher
316
317            self._iteration += 1
318            n_iter += 1
319            if self._iteration >= self.max_iteration:
320                break
321            progress.update(1)
322
323        t_per_iter = (time.time() - t_per_iter) / n_iter
324        return t_per_iter
325
326    def _validate_supervised(self, forward_context):
327        metric_val = 0.0
328        loss_val = 0.0
329
330        for x, y in self.supervised_val_loader:
331            x, y = (
332                x.to(self.device, non_blocking=True),
333                y.to(self.device, non_blocking=True)
334            )
335
336            with forward_context():
337                pred = self.model(x)
338                loss, metric = self.supervised_loss_and_metric(pred, y)
339                loss_val += loss.item()
340            metric_val += metric.item()
341
342        metric_val /= len(self.supervised_val_loader)
343        loss_val /= len(self.supervised_val_loader)
344
345        if self.logger is not None:
346            self.logger.log_validation_supervised(
347                self._iteration, metric_val, loss_val, x, y, pred
348            )
349
350        return metric_val
351
352    def _validate_unsupervised(self, forward_context):
353        metric_val = 0.0
354        loss_val = 0.0
355
356        for x in self.unsupervised_val_loader:
357            self.augmenter.reset_all()
358            x = x.to(self.device, non_blocking=True)
359
360            # apply augmentations
361            x_w = self.augmenter.weak.transform(x)
362            x_s1, x_s2 = self.augmenter.strong1.transform(x), self.augmenter.strong2.transform(x)
363
364            # Compute the pseudo labels (unsupervised teacher prediction)
365            with forward_context():
366                pseudo_labels, label_filter = self.pseudo_labeler(self.teacher, x_w)
367                pseudo_labels_inv = self.augmenter.weak.reverse_transform(pseudo_labels)
368                label_filter_inv = (
369                    self.augmenter.weak.reverse_transform(label_filter)
370                    if label_filter is not None else None
371                )
372
373                if self.complementary_dropout:
374                    pred_s1, pred_s2 = self.predict_with_comp_drop(self.model, torch.cat((x_s1, x_s2))).chunk(2)
375                else:
376                    pred_s1, pred_s2 = self.model(torch.cat((x_s1, x_s2))).chunk(2)
377                pred_s1_inv = self.augmenter.strong1.reverse_transform(pred_s1)
378                pred_s2_inv = self.augmenter.strong2.reverse_transform(pred_s2)
379
380                loss, metric = self.unsupervised_loss_and_metric(
381                    torch.stack((pred_s1_inv, pred_s2_inv)),
382                    pseudo_labels_inv,
383                    label_filter_inv,
384                    pred_dim=2,
385                )
386            loss_val += loss.item()
387            metric_val += metric.item()
388
389        metric_val /= len(self.unsupervised_val_loader)
390        loss_val /= len(self.unsupervised_val_loader)
391
392        if self.logger is not None:
393            self.logger.log_validation_unsupervised(
394                self._iteration,
395                metric_val,
396                loss_val,
397                x,
398                pred_s1_inv,
399                pred_s2_inv,
400                pseudo_labels_inv,
401                label_filter_inv,
402            )
403
404            self.logger.log_validation_augmentations(
405                self._iteration,
406                x_w,
407                x_s1,
408                x_s2,
409                pseudo_labels,
410                pred_s1,
411                pred_s2,
412            )
413
414        self.pseudo_labeler.step(metric_val, self._epoch)
415
416        return metric_val
417
418    def _validate_impl(self, forward_context):
419        self.model.eval()
420
421        with torch.no_grad():
422
423            if self.supervised_val_loader is None:
424                supervised_metric = None
425            else:
426                supervised_metric = self._validate_supervised(forward_context)
427
428            if self.unsupervised_val_loader is None:
429                unsupervised_metric = None
430            else:
431                unsupervised_metric = self._validate_unsupervised(forward_context)
432
433        if unsupervised_metric is None:
434            metric = supervised_metric
435        elif supervised_metric is None:
436            metric = unsupervised_metric
437        else:
438            metric = (supervised_metric + unsupervised_metric) / 2
439
440        return metric