clarena.cl_datasets.split_mnist

  1r"""
  2The submodule in `cl_datasets` for Split MNIST dataset.
  3"""
  4
  5__all__ = ["SplitMNIST"]
  6
  7import logging
  8from typing import Callable
  9
 10import torch
 11from torch.utils.data import Dataset, random_split
 12from torchvision.datasets import MNIST
 13from torchvision.transforms import transforms
 14
 15from clarena.cl_datasets import CLSplitDataset
 16
 17# always get logger for built-in logging in each module
 18pylogger = logging.getLogger(__name__)
 19
 20
 21class SplitMNIST(CLSplitDataset):
 22    r"""Split MNIST dataset. The [MNIST dataset](http://yann.lecun.com/exdb/mnist/) is a collection of handwritten digits. It consists of 60,000 training and 10,000 test images of handwritten digit images (10 classes), each 28x28 grayscale image."""
 23
 24    original_dataset_python_class: type[Dataset] = MNIST
 25    r"""The original dataset class."""
 26
 27    def __init__(
 28        self,
 29        root: str,
 30        class_split: dict[int, list[int]],
 31        validation_percentage: float,
 32        batch_size: int | dict[int, int] = 1,
 33        num_workers: int | dict[int, int] = 0,
 34        custom_transforms: (
 35            Callable
 36            | transforms.Compose
 37            | None
 38            | dict[int, Callable | transforms.Compose | None]
 39        ) = None,
 40        repeat_channels: int | None | dict[int, int | None] = None,
 41        to_tensor: bool | dict[int, bool] = True,
 42        resize: tuple[int, int] | None | dict[int, tuple[int, int] | None] = None,
 43    ) -> None:
 44        r"""
 45        **Args:**
 46        - **root** (`str`): the root directory where the original MNIST data 'MNIST/' live.
 47        - **class_split** (`dict[int, list[int]]`): the dict of classes for each task. The keys are task IDs ane the values are lists of class labels (integers starting from 0) to split for each task.
 48        - **validation_percentage** (`float`): The percentage to randomly split some training data into validation data.
 49        - **batch_size** (`int` | `dict[int, int]`): the batch size for train, val, and test dataloaders.
 50        If it is a dict, the keys are task IDs and the values are the batch sizes for each task. If it is an `int`, it is the same batch size for all tasks.
 51        - **num_workers** (`int` | `dict[int, int]`): the number of workers for dataloaders.
 52        If it is a dict, the keys are task IDs and the values are the number of workers for each task. If it is an `int`, it is the same number of workers for all tasks.
 53        - **custom_transforms** (`transform` or `transforms.Compose` or `None` or dict of them): the custom transforms to apply ONLY to the TRAIN dataset. Can be a single transform, composed transforms, or no transform. `ToTensor()`, normalization, permute, and so on are not included.
 54        If it is a dict, the keys are task IDs and the values are the custom transforms for each task. If it is a single transform or composed transforms, it is applied to all tasks. If it is `None`, no custom transforms are applied.
 55        - **repeat_channels** (`int` | `None` | dict of them): the number of channels to repeat for each task. Default is `None`, which means no repeat.
 56        If it is a dict, the keys are task IDs and the values are the number of channels to repeat for each task. If it is an `int`, it is the same number of channels to repeat for all tasks. If it is `None`, no repeat is applied.
 57        - **to_tensor** (`bool` | `dict[int, bool]`): whether to include the `ToTensor()` transform. Default is `True`.
 58        If it is a dict, the keys are task IDs and the values are whether to include the `ToTensor()` transform for each task. If it is a single boolean value, it is applied to all tasks.
 59        - **resize** (`tuple[int, int]` | `None` or dict of them): the size to resize the images to. Default is `None`, which means no resize.
 60        If it is a dict, the keys are task IDs and the values are the sizes to resize for each task. If it is a single tuple of two integers, it is applied to all tasks. If it is `None`, no resize is applied.
 61        """
 62
 63        super().__init__(
 64            root=root,
 65            class_split=class_split,
 66            batch_size=batch_size,
 67            num_workers=num_workers,
 68            custom_transforms=custom_transforms,
 69            repeat_channels=repeat_channels,
 70            to_tensor=to_tensor,
 71            resize=resize,
 72        )
 73
 74        self.validation_percentage: float = validation_percentage
 75        r"""The percentage to randomly split some training data into validation data."""
 76
 77    def prepare_data(self) -> None:
 78        r"""Download the original MNIST dataset if haven't."""
 79
 80        if self.task_id != 1:
 81            return  # download all original datasets only at the beginning of first task
 82
 83        MNIST(root=self.root_t, train=True, download=True)
 84        MNIST(root=self.root_t, train=False, download=True)
 85
 86        pylogger.debug(
 87            "The original MNIST dataset has been downloaded to %s.", self.root
 88        )
 89
 90    def get_subset_of_classes(self, dataset: Dataset) -> Dataset:
 91        r"""Get a subset of classes from the dataset of current classes of `self.task_id`. It is used when constructing the split. It must be implemented by subclasses.
 92
 93        **Args:**
 94        - **dataset** (`Dataset`): the dataset to retrieve subset from.
 95
 96        **Returns:**
 97        - **subset** (`Dataset`): the subset of classes from the dataset.
 98        """
 99        classes = self.class_split[self.task_id]
100
101        # get the indices of the dataset that belong to the classes
102        idx = [i for i, (_, target) in enumerate(dataset) if target in classes]
103
104        # subset the dataset by the indices, in-place operation
105        dataset.data = dataset.data[idx]  # data is a Numpy ndarray
106        dataset.targets = [dataset.targets[i] for i in idx]  # targets is a list
107
108        dataset.target_transform = (
109            self.target_transform()
110        )  # cl class mapping should be applied after the split
111
112        return dataset
113
114    def train_and_val_dataset(self) -> tuple[Dataset, Dataset]:
115        r"""Get the training and validation dataset of task `self.task_id`.
116
117        **Returns:**
118        - **train_and_val_dataset** (`tuple[Dataset, Dataset]`): the train and validation dataset of task `self.task_id`.
119        """
120        dataset_train_and_val = self.get_subset_of_classes(
121            MNIST(
122                root=self.root_t,
123                train=True,
124                transform=self.train_and_val_transforms(),
125                # cl class mapping should be applied after the split
126                download=False,
127            )
128        )
129
130        return random_split(
131            dataset_train_and_val,
132            lengths=[1 - self.validation_percentage, self.validation_percentage],
133            generator=torch.Generator().manual_seed(
134                42
135            ),  # this must be set fixed to make sure the datasets across experiments are the same. Don't handle it to global seed as it might vary across experiments
136        )
137
138    def test_dataset(self) -> Dataset:
139        r"""Get the test dataset of task `self.task_id`.
140
141        **Returns:**
142        - **test_dataset** (`Dataset`): the test dataset of task `self.task_id`.
143        """
144        dataset_test = self.get_subset_of_classes(
145            MNIST(
146                root=self.root_t,
147                train=False,
148                transform=self.test_transforms(),
149                # cl class mapping should be applied after the split
150                download=False,
151            )
152        )
153
154        return dataset_test