torch_em.self_training.logger

  1import os
  2
  3import torch_em
  4import torch
  5
  6from torchvision.utils import make_grid
  7from torch.utils.tensorboard import SummaryWriter
  8
  9
 10class SelfTrainingTensorboardLogger(torch_em.trainer.logger_base.TorchEmLogger):
 11    """Logger for self-training via `torch_em.self_training.FixMatch` or `torch_em.self_training.MeanTeacher`.
 12    Also supports logging training with invertible augmentations.
 13
 14    Args:
 15        trainer: The instantiated trainer class.
 16        save_root: The root directory for saving the checkpoints and logs.
 17    """
 18    @staticmethod
 19    def _get_image_channel(x):
 20        return x[0, 0:1] if x.shape[1] > 1 else x[0]
 21
 22    def __init__(self, trainer, save_root, **unused_kwargs):
 23        super().__init__(trainer, save_root)
 24        self.my_root = save_root
 25        self.log_dir = f"./logs/{trainer.name}" if self.my_root is None else\
 26            os.path.join(self.my_root, "logs", trainer.name)
 27        os.makedirs(self.log_dir, exist_ok=True)
 28
 29        self.tb = SummaryWriter(self.log_dir)
 30        self.log_image_interval = trainer.log_image_interval
 31
 32    def _add_supervised_images(self, step, name, x, y, pred):
 33        if x.ndim == 5:
 34            assert y.ndim == pred.ndim == 5
 35            zindex = x.shape[2] // 2
 36            x, y, pred = x[:, :, zindex], y[:, :, zindex], pred[:, :, zindex]
 37
 38        num_channels = y.shape[1]
 39
 40        images = (
 41            [torch_em.transform.raw.normalize(self._get_image_channel(x))] * num_channels +
 42            [y[0, c:c+1] for c in range(num_channels)] +
 43            [pred[0, c:c+1] for c in range(num_channels)]
 44        )
 45        grid = make_grid(images, nrow=num_channels, padding=8)
 46        self.tb.add_image(tag=f"{name}/supervised/input-labels-prediction", img_tensor=grid, global_step=step)
 47
 48    def _add_unsupervised_images(self, step, name, x1, x2, pred, pseudo_labels, label_filter):
 49        if x1.ndim == 5:
 50            assert x2.ndim == pred.ndim == pseudo_labels.ndim == 5
 51            zindex = x1.shape[2] // 2
 52            x1, x2, pred = x1[:, :, zindex], x2[:, :, zindex], pred[:, :, zindex]
 53            pseudo_labels = pseudo_labels[:, :, zindex]
 54            if label_filter is not None:
 55                assert label_filter.ndim == 5
 56                label_filter = label_filter[:, :, zindex]
 57
 58        num_channels = pred.shape[1]
 59
 60        images = (
 61            [torch_em.transform.raw.normalize(self._get_image_channel(x1))] +
 62            [torch_em.transform.raw.normalize(self._get_image_channel(x2))] +
 63            [torch.zeros_like(self._get_image_channel(x1))] * (num_channels - 2) +
 64            [pred[0, c:c+1] for c in range(num_channels)] +
 65            [pseudo_labels[0, c:c+1] for c in range(num_channels)]
 66        )
 67        im_name = f"{name}/unsupervised/image-prediction-pseudolabels"
 68        # if trainer with invertible augmentations, untransformed images
 69        # and inverted pred/labels are logged for better visual comparison,
 70        # otherwise the transformed images are logged
 71        if label_filter is not None:
 72            images.extend([label_filter[0, c:c+1] for c in range(num_channels)])
 73            im_name += "-labelfilter"
 74        grid = make_grid(images, nrow=num_channels, padding=8)
 75        self.tb.add_image(tag=im_name, img_tensor=grid, global_step=step)
 76
 77    def log_combined_loss(self, step, loss):
 78        """@private
 79        """
 80        self.tb.add_scalar(tag="train/combined_loss", scalar_value=loss, global_step=step)
 81
 82    def log_lr(self, step, lr):
 83        """@private
 84        """
 85        self.tb.add_scalar(tag="train/learning_rate", scalar_value=lr, global_step=step)
 86
 87    def log_train_supervised(self, step, loss, x, y, pred):
 88        """@private
 89        """
 90        self.tb.add_scalar(tag="train/supervised/loss", scalar_value=loss, global_step=step)
 91        if step % self.log_image_interval == 0:
 92            self._add_supervised_images(step, "train", x, y, pred)
 93
 94    def log_validation_supervised(self, step, metric, loss, x, y, pred):
 95        """@private
 96        """
 97        self.tb.add_scalar(tag="validation/supervised/loss", scalar_value=loss, global_step=step)
 98        self.tb.add_scalar(tag="validation/supervised/metric", scalar_value=metric, global_step=step)
 99        self._add_supervised_images(step, "validation", x, y, pred)
100
101    def log_train_unsupervised(self, step, loss, x1, x2, pred, pseudo_labels, label_filter=None):
102        """@private
103        """
104        self.tb.add_scalar(tag="train/unsupervised/loss", scalar_value=loss, global_step=step)
105        if step % self.log_image_interval == 0:
106            self._add_unsupervised_images(step, "train", x1, x2, pred, pseudo_labels, label_filter)
107
108    def log_validation_unsupervised(self, step, metric, loss, x1, x2, pred, pseudo_labels, label_filter=None):
109        """@private
110        """
111        self.tb.add_scalar(tag="validation/unsupervised/loss", scalar_value=loss, global_step=step)
112        self.tb.add_scalar(tag="validation/unsupervised/metric", scalar_value=metric, global_step=step)
113        self._add_unsupervised_images(step, "validation", x1, x2, pred, pseudo_labels, label_filter)
114
115    def log_validation(self, step, metric, loss, gt_metric=None):
116        """@private
117        """
118        self.tb.add_scalar(tag="validation/loss", scalar_value=loss, global_step=step)
119        self.tb.add_scalar(tag="validation/metric", scalar_value=metric, global_step=step)
120        if gt_metric is not None:
121            self.tb.add_scalar(tag="validation/gt_metric", scalar_value=gt_metric, global_step=step)
122
123    def log_ct(self, step, ct):
124        self.tb.add_scalar(tag="train/confidence_threshold", scalar_value=ct, global_step=step)
125
126    def _add_augmented_images(self, step, name, xu1, xu2, pseudo_labels, pred):
127        if xu1.ndim == 5:
128            assert xu2.ndim == pseudo_labels.ndim == pred.ndim == 5
129            zindex = xu1.shape[2] // 2
130            xu1 = xu1[:, :, zindex]
131            xu2 = xu2[:, :, zindex]
132            pred = pred[:, :, zindex]
133            pseudo_labels = pseudo_labels[:, :, zindex]
134
135        images = [
136            torch_em.transform.raw.normalize(xu1[0]),
137            torch_em.transform.raw.normalize(xu2[0]),
138            pseudo_labels[0, 0:1],
139            pred[0, 0:1],
140        ]
141        grid = make_grid(images, nrow=2, padding=8)
142        self.tb.add_image(tag=f"{name}/unsupervised/aug1-aug2-pseudolabels-prediction", img_tensor=grid, global_step=step)
143
144    def log_train_augmentations(self, step, xu1, xu2, pseudo_labels, pred):
145        if step % self.log_image_interval == 0:
146            self._add_augmented_images(step, "train_augmentations", xu1, xu2, pseudo_labels, pred)
147
148    def log_validation_augmentations(self, step, xu1, xu2, pseudo_labels, pred):
149        if step % self.log_image_interval == 0:
150            self._add_augmented_images(step, "validation_augmentations", xu1, xu2, pseudo_labels, pred)
151
152
153class ProbabilisticUNetTrainerLogger(torch_em.trainer.logger_base.TorchEmLogger):
154    def __init__(self, trainer, save_root, **unused_kwargs):
155        super().__init__(trainer, save_root)
156        self.log_dir = f"./logs/{trainer.name}" if save_root is None else\
157            os.path.join(save_root, "logs", trainer.name)
158        os.makedirs(self.log_dir, exist_ok=True)
159
160        self.tb = torch.utils.tensorboard.SummaryWriter(self.log_dir)
161        self.log_image_interval = trainer.log_image_interval
162
163    def add_image(self, x, y, samples, name, step):
164        # NOTE: we only show the first tensor per batch for all images
165        self.tb.add_image(tag=f"{name}/input", img_tensor=x[0], global_step=step)
166        target_grid = make_grid(list(y[0].float().split(1)), nrow=4, padding=4, normalize=True)
167        self.tb.add_image(tag=f"{name}/target", img_tensor=target_grid, global_step=step)
168        predictions = [
169            sample.sigmoid() if sample.shape[1] == 1 else sample.argmax(dim=1, keepdim=True) for sample in samples
170        ]
171        sample_grid = make_grid([sample[0].float() for sample in predictions], nrow=4, padding=4, normalize=True)
172        self.tb.add_image(tag=f"{name}/samples", img_tensor=sample_grid, global_step=step)
173
174    def log_train(self, step, loss, lr, x, y, samples):
175        self.tb.add_scalar(tag="train/loss", scalar_value=loss, global_step=step)
176        self.tb.add_scalar(tag="train/learning_rate", scalar_value=lr, global_step=step)
177        if step % self.log_image_interval == 0:
178            self.add_image(x, y, samples, "train", step)
179
180    def log_validation(self, step, metric, loss, x, y, samples):
181        self.tb.add_scalar(tag="validation/loss", scalar_value=loss, global_step=step)
182        self.tb.add_scalar(tag="validation/metric", scalar_value=metric, global_step=step)
183        self.add_image(x, y, samples, "validation", step)
184
185
186class UniMatchv2TensorboardLogger(torch_em.trainer.logger_base.TorchEmLogger):
187    """Logger for self-training via `torch_em.self_training.UniMatchv2Trainer`.
188
189    Args:
190        trainer: The instantiated trainer class.
191        save_root: The root directory for saving the checkpoints and logs.
192    """
193
194    def __init__(self, trainer, save_root, **unused_kwargs):
195        super().__init__(trainer, save_root)
196        self.my_root = save_root
197        self.log_dir = (
198            f"./logs/{trainer.name}"
199            if self.my_root is None
200            else os.path.join(self.my_root, "logs", trainer.name)
201        )
202        os.makedirs(self.log_dir, exist_ok=True)
203
204        self.tb = SummaryWriter(self.log_dir)
205        self.log_image_interval = trainer.log_image_interval
206
207    def _add_supervised_images(self, step, name, x, y, pred):
208        if x.ndim == 5:
209            assert y.ndim == pred.ndim == 5
210            zindex = x.shape[2] // 2
211            x, y, pred = x[:, :, zindex], y[:, :, zindex], pred[:, :, zindex]
212
213        num_channels = y.shape[1]
214
215        images = (
216            [torch_em.transform.raw.normalize(x[0])] * num_channels +
217            [y[0, c:c+1] for c in range(num_channels)] +
218            [pred[0, c:c+1] for c in range(num_channels)]
219        )
220        grid = make_grid(images, nrow=num_channels, padding=8)
221        self.tb.add_image(
222            tag=f"{name}/supervised/input-labels-prediction",
223            img_tensor=grid,
224            global_step=step,
225        )
226
227    def _add_unsupervised_images(
228        self, step, name, x, pred_s1, pred_s2, pseudo_labels, label_filter
229    ):
230        if x.ndim == 5:
231            assert (
232                pred_s1.ndim
233                == pred_s2.ndim
234                == pseudo_labels.ndim
235                == 5
236            )
237            zindex = x.shape[2] // 2
238            x = x[:, :, zindex]
239            pred_s1, pred_s2 = pred_s1[:, :, zindex], pred_s2[:, :, zindex]
240            pseudo_labels = pseudo_labels[:, :, zindex]
241            if label_filter is not None:
242                assert label_filter.ndim == 5
243                label_filter = label_filter[:, :, zindex]
244        num_channels = pred_s1.shape[1]
245
246        images = (
247            [torch_em.transform.raw.normalize(x[0])] * num_channels +
248            [pred_s1[0, c:c+1] for c in range(num_channels)] +
249            [pred_s2[0, c:c+1] for c in range(num_channels)] +
250            [pseudo_labels[0, c:c+1] for c in range(num_channels)]
251        )
252
253        im_name = (
254            f"{name}/unsupervised/image-pred_s1-pred_s2-pseudolabels"
255        )
256        if label_filter is not None:
257            images.extend([label_filter[0, c:c+1] for c in range(num_channels)])
258            im_name += "-labelfilter"
259        grid = make_grid(images, nrow=num_channels, padding=8)
260        self.tb.add_image(tag=im_name, img_tensor=grid, global_step=step)
261
262    def log_combined_loss(self, step, loss):
263        """@private"""
264        self.tb.add_scalar(
265            tag="train/combined_loss", scalar_value=loss, global_step=step
266        )
267
268    def log_lr(self, step, lr):
269        """@private"""
270        self.tb.add_scalar(tag="train/learning_rate", scalar_value=lr, global_step=step)
271
272    def log_train_supervised(self, step, loss, x, y, pred):
273        """@private"""
274        self.tb.add_scalar(
275            tag="train/supervised/loss", scalar_value=loss, global_step=step
276        )
277        if step % self.log_image_interval == 0:
278            self._add_supervised_images(step, "train", x, y, pred)
279
280    def log_validation_supervised(self, step, metric, loss, x, y, pred):
281        """@private"""
282        self.tb.add_scalar(
283            tag="validation/supervised/loss", scalar_value=loss, global_step=step
284        )
285        self.tb.add_scalar(
286            tag="validation/supervised/metric", scalar_value=metric, global_step=step
287        )
288        self._add_supervised_images(step, "validation", x, y, pred)
289
290    def log_train_unsupervised(
291        self,
292        step,
293        loss,
294        x,
295        pred_s1,
296        pred_s2,
297        pseudo_labels,
298        label_filter=None,
299    ):
300        """@private"""
301        self.tb.add_scalar(
302            tag="train/unsupervised/loss", scalar_value=loss, global_step=step
303        )
304        if step % self.log_image_interval == 0:
305            self._add_unsupervised_images(
306                step,
307                "train",
308                x,
309                pred_s1,
310                pred_s2,
311                pseudo_labels,
312                label_filter,
313            )
314
315    def log_validation_unsupervised(
316        self,
317        step,
318        metric,
319        loss,
320        x,
321        pred_s1,
322        pred_s2,
323        pseudo_labels,
324        label_filter=None,
325    ):
326        """@private"""
327        self.tb.add_scalar(
328            tag="validation/unsupervised/loss", scalar_value=loss, global_step=step
329        )
330        self.tb.add_scalar(
331            tag="validation/unsupervised/metric", scalar_value=metric, global_step=step
332        )
333        self._add_unsupervised_images(
334            step,
335            "validation",
336            x,
337            pred_s1,
338            pred_s2,
339            pseudo_labels,
340            label_filter,
341        )
342
343    def log_ct(self, step, ct):
344        self.tb.add_scalar(
345            tag="train/confidence_threshold", scalar_value=ct, global_step=step
346        )
347
348    # LOG AUGMENTATIONS FOR DEBUGGING ###
349    def _add_augmented_images(
350        self, step, name, x_u_w, x_u_s1, x_u_s2, pseudo_labels, pred_s1, pred_s2
351    ):
352        if x_u_w.ndim == 5:
353            assert (
354                x_u_s1.ndim
355                == x_u_s2.ndim
356                == pseudo_labels.ndim
357                == pred_s1.ndim
358                == pred_s2.ndim
359                == 5
360            )
361            zindex = x_u_w.shape[2] // 2
362            x_u_w = x_u_w[:, :, zindex]
363            x_u_s1, x_u_s2 = x_u_s1[:, :, zindex], x_u_s2[:, :, zindex]
364            pred_s1, pred_s2 = pred_s1[:, :, zindex], pred_s2[:, :, zindex]
365            pseudo_labels = pseudo_labels[:, :, zindex]
366
367        images = [
368            torch_em.transform.raw.normalize(x_u_w[0]),
369            torch_em.transform.raw.normalize(x_u_s1[0]),
370            torch_em.transform.raw.normalize(x_u_s2[0]),
371            pseudo_labels[0, 0:1],
372            pred_s1[0, 0:1],
373            pred_s2[0, 0:1],
374        ]
375        im_name = (
376            f"{name}/unsupervised/aug_w-aug_s1-aug_s2-pseudolabels-pred_s1-pred_s2"
377        )
378        grid = make_grid(images, nrow=3, padding=8)
379        self.tb.add_image(tag=im_name, img_tensor=grid, global_step=step)
380
381    def log_train_augmentations(
382        self, step, x_u_w, x_u_s1, x_u_s2, pseudo_labels, pred_s1, pred_s2
383    ):
384        if step % self.log_image_interval == 0:
385            self._add_augmented_images(
386                step,
387                "train_augmentations",
388                x_u_w,
389                x_u_s1,
390                x_u_s2,
391                pseudo_labels,
392                pred_s1,
393                pred_s2,
394            )
395
396    def log_validation_augmentations(
397        self, step, x_u_w, x_u_s1, x_u_s2, pseudo_labels, pred_s1, pred_s2
398    ):
399        if step % self.log_image_interval == 0:
400            self._add_augmented_images(
401                step,
402                "validation_augmentations",
403                x_u_w,
404                x_u_s1,
405                x_u_s2,
406                pseudo_labels,
407                pred_s1,
408                pred_s2,
409            )