clarena.callbacks.save_first_batch_images

  1r"""
  2The submodule in `callbacks` for callback of saving first batch images.
  3"""
  4
  5__all__ = ["SaveFirstBatchImages"]
  6
  7
  8import logging
  9import os
 10
 11import torch
 12import torchvision
 13import wandb
 14from lightning import Callback, Trainer
 15from lightning.pytorch.loggers import WandbLogger
 16
 17from clarena.cl_algorithms import CLAlgorithm
 18from clarena.mtl_algorithms import MTLAlgorithm
 19from clarena.stl_algorithms import STLAlgorithm
 20
 21# always get logger for built-in logging in each module
 22pylogger = logging.getLogger(__name__)
 23
 24
 25class SaveFirstBatchImages(Callback):
 26    r"""Saves images and labels of the first batch of training data into files. In continual learning / unlearning, applies to all tasks."""
 27
 28    def __init__(
 29        self,
 30        save_dir: str,
 31        img_prefix: str = "sample",
 32        labels_filename: str = "labels.txt",
 33        task_ids_filename: str | None = "tasks.txt",
 34    ) -> None:
 35        r"""Initialize the Save First Batch Images Callback.
 36
 37        **Args:**
 38        - **save_dir** (`str`): the directory to save images and labels as files. Better inside the output directory.
 39        - **img_prefix** (`str`): the prefix for image files.
 40        - **labels_filename** (`str`): the filename for the labels file as texts.
 41        - **task_ids_filename** (`str` | `None`): the filename for the task IDs file as texts. Only used in MTL algorithms. If `None`, no task IDs file is saved.
 42        """
 43
 44        os.makedirs(save_dir, exist_ok=True)
 45
 46        self.save_dir: str = save_dir
 47        r"""Store the directory to save images and labels as files."""
 48        self.img_prefix: str = img_prefix
 49        r"""Store the prefix for image files."""
 50        self.labels_filename: str = labels_filename
 51        r"""Store the filename for the labels file."""
 52        self.task_ids_filename: str | None = task_ids_filename
 53        r"""Store the filename for the task IDs file. Only used in MTL algorithms. If `None`, no task IDs file is saved."""
 54
 55        self.called: bool = False
 56        r"""Flag to avoid calling the callback multiple times."""
 57
 58    def on_train_batch_end(
 59        self,
 60        trainer: Trainer,
 61        pl_module: CLAlgorithm | MTLAlgorithm | STLAlgorithm,
 62        outputs,
 63        batch,
 64        batch_idx: int,
 65        dataloader_idx: int = 0,
 66    ) -> None:
 67        r"""Save images and labels into files in the first batch of training data at the beginning of the training of the task."""
 68
 69        if isinstance(pl_module, CLAlgorithm):
 70            image_batch, label_batch = batch
 71            image_samples = list(torch.unbind(image_batch, dim=0))
 72            label_samples = list(torch.unbind(label_batch, dim=0))
 73
 74            # save images and labels as documents
 75            save_dir_task = os.path.join(self.save_dir, f"task_{pl_module.task_id}")
 76            os.makedirs(save_dir_task, exist_ok=True)
 77            labels_file = open(
 78                os.path.join(save_dir_task, self.labels_filename),
 79                "w",
 80                encoding="utf-8",
 81            )
 82            for i, (image, label) in enumerate(zip(image_samples, label_samples)):
 83                torchvision.utils.save_image(
 84                    image, os.path.join(save_dir_task, f"{self.img_prefix}_{i}.png")
 85                )
 86                labels_file.write(f"{self.img_prefix}_{i}.png: {label}\n")
 87            labels_file.close()
 88        elif isinstance(pl_module, MTLAlgorithm):
 89            image_batch, label_batch, tasks_batch = batch
 90            image_samples = list(torch.unbind(image_batch, dim=0))
 91            label_samples = list(torch.unbind(label_batch, dim=0))
 92            task_samples = list(torch.unbind(tasks_batch, dim=0))
 93
 94            # save images, labels and task_ids as documents
 95            labels_file = open(
 96                os.path.join(self.save_dir, self.labels_filename), "w", encoding="utf-8"
 97            )
 98            task_ids_file = open(
 99                os.path.join(self.save_dir, self.task_ids_filename),
100                "w",
101                encoding="utf-8",
102            )
103            for i, (image, label, task_id) in enumerate(
104                zip(image_samples, label_samples, task_samples)
105            ):
106                torchvision.utils.save_image(
107                    image, os.path.join(self.save_dir, f"{self.img_prefix}_{i}.png")
108                )
109                labels_file.write(f"{self.img_prefix}_{i}.png: {label}\n")
110                task_ids_file.write(f"{self.img_prefix}_{i}.png: {task_id}\n")
111            labels_file.close()
112            task_ids_file.close()
113
114        elif isinstance(pl_module, STLAlgorithm):
115            image_batch, label_batch = batch
116            image_samples = list(torch.unbind(image_batch, dim=0))
117            label_samples = list(torch.unbind(label_batch, dim=0))
118
119            # save images and labels as documents
120            labels_file = open(
121                os.path.join(self.save_dir, self.labels_filename), "w", encoding="utf-8"
122            )
123            for i, (image, label) in enumerate(zip(image_samples, label_samples)):
124                torchvision.utils.save_image(
125                    image, os.path.join(self.save_dir, f"{self.img_prefix}_{i}.png")
126                )
127                labels_file.write(f"{self.img_prefix}_{i}.png: {label}\n")
128            labels_file.close()
129
130        self.called = True  # flag to avoid calling the callback multiple times
131
132    def on_validation_batch_end(
133        self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx: int = 0
134    ):
135        """Called when the validation batch ends."""
136
137        wandb_logger = wandb_logger = next(
138            (logger for logger in pl_module.loggers if isinstance(logger, WandbLogger)),
139            None,
140        )
141
142        # `outputs` comes from `LightningModule.validation_step`
143        # which corresponds to our model predictions in this case
144
145        # Let's log 20 sample image predictions from the first batch
146        if batch_idx == 0:
147            n = 20
148            x, y = batch[:2]
149            x = (x - x.min()) / (x.max() - x.min())
150            images = [wandb.Image(img) for img in x[:n]]
151            captions = [
152                f"Ground Truth: {y_i} - Prediction: {y_pred}"
153                for y_i, y_pred in zip(y[:n], outputs["preds"][:n])
154            ]
155
156            # Option 1: log images with `WandbLogger.log_image`
157            if wandb_logger is not None:
158                wandb_logger.log_image(
159                    key="sample_images", images=images, caption=captions
160                )
161
162            # Option 2: log images and predictions as a W&B Table
163            columns = ["image", "ground truth", "prediction"]
164            data = [
165                [wandb.Image(x_i), y_i, y_pred]
166                for x_i, y_i, y_pred in list(zip(x[:n], y[:n], outputs["preds"][:n]))
167            ]
168            if wandb_logger is not None:
169                wandb_logger.log_table(key="sample_table", columns=columns, data=data)