clarena.cl_datasets.permuted_mnist

  1r"""
  2The submodule in `cl_datasets` for Permuted MNIST dataset.
  3"""
  4
  5__all__ = ["PermutedMNIST"]
  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 CLPermutedDataset
 16
 17# always get logger for built-in logging in each module
 18pylogger = logging.getLogger(__name__)
 19
 20
 21class PermutedMNIST(CLPermutedDataset):
 22    r"""Permuted MNIST dataset. The [original MNIST dataset](http://yann.lecun.com/exdb/mnist/) is a collection of handwritten digits. It consists of 70,000 28x28 B&W images in 10 classes (correspond to 10 digits), with 7000 images per class. There are 60,000 training examples and 10,000 test examples."""
 23
 24    num_classes: int = 10
 25    """The number of classes in MNIST."""
 26
 27    img_size: torch.Size = torch.Size([1, 28, 28])
 28    """The size of MNIST images."""
 29
 30    mean_original: tuple[float] = (0.1307,)
 31    """The mean values for normalisation."""
 32
 33    std_original: tuple[float] = (0.3081,)
 34    """The standard deviatfion values for normalisation."""
 35
 36    def __init__(
 37        self,
 38        root: str,
 39        num_tasks: int,
 40        validation_percentage: float,
 41        batch_size: int = 1,
 42        num_workers: int = 8,
 43        custom_transforms: Callable | transforms.Compose | None = None,
 44        custom_target_transforms: Callable | transforms.Compose | None = None,
 45        permutation_mode: str = "first_channel_only",
 46        permutation_seeds: list[int] | None = None,
 47    ) -> None:
 48        r"""Initialise the Permuted MNIST dataset.
 49
 50        **Args:**
 51        - **root** (`str`): the root directory where the original MNIST data 'MNIST/raw/train-images-idx3-ubyte' and 'MNIST/raw/t10k-images-idx3-ubyte' live.
 52        - **num_tasks** (`int`): the maximum number of tasks supported by the CL dataset.
 53        - **validation_percentage** (`float`): the percentage to randomly split some of the training data into validation data.
 54        - **batch_size** (`int`): The batch size in train, val, test dataloader.
 55        - **num_workers** (`int`): the number of workers for dataloaders.
 56        - **custom_transforms** (`transform` or `transforms.Compose` or `None`): the custom transforms to apply to ONLY TRAIN dataset. Can be a single transform, composed transforms or no transform.
 57        `ToTensor()`, normalise, permute and so on are not included.
 58        - **custom_target_transforms** (`transform` or `transforms.Compose` or `None`): the custom target transforms to apply to dataset labels. Can be a single transform, composed transforms or no transform. CL class mapping is not included.
 59        - **permutation_mode** (`str`): the mode of permutation, should be one of the following:
 60            1. 'all': permute all pixels.
 61            2. 'by_channel': permute channel by channel separately. All channels are applied the same permutation order.
 62            3. 'first_channel_only': permute only the first channel.
 63        - **permutation_seeds** (`list[int]` or `None`): the seeds for permutation operations used to construct tasks. Make sure it has the same number of seeds as `num_tasks`. Default is None, which creates a list of seeds from 1 to `num_tasks`.
 64        """
 65        CLPermutedDataset.__init__(
 66            self,
 67            root=root,
 68            num_tasks=num_tasks,
 69            validation_percentage=validation_percentage,
 70            batch_size=batch_size,
 71            num_workers=num_workers,
 72            custom_transforms=custom_transforms,
 73            custom_target_transforms=custom_target_transforms,
 74            permutation_mode=permutation_mode,
 75            permutation_seeds=permutation_seeds,
 76        )
 77
 78    def prepare_data(self) -> None:
 79        r"""Download the original MNIST dataset if haven't."""
 80        # just download
 81        MNIST(root=self.root, train=True, download=True)
 82        MNIST(root=self.root, train=False, download=True)
 83
 84        pylogger.debug(
 85            "The original MNIST dataset has been downloaded to %s.", self.root
 86        )
 87
 88    def train_and_val_dataset(self) -> tuple[Dataset, Dataset]:
 89        """Get the training and validation dataset of task `self.task_id`.
 90
 91        **Returns:**
 92        - **train_and_val_dataset** (`tuple[Dataset, Dataset]`): the train and validation dataset of task `self.task_id`.
 93        """
 94        dataset_train_and_val = MNIST(
 95            root=self.root,
 96            train=True,
 97            transform=self.train_and_val_transforms(to_tensor=True),
 98            download=False,
 99        )
100        return random_split(
101            dataset_train_and_val,
102            lengths=[1 - self.validation_percentage, self.validation_percentage],
103        )
104
105    def test_dataset(self) -> Dataset:
106        r"""Get the test dataset of task `self.task_id`.
107
108        **Returns:**
109        - **test_dataset** (`Dataset`): the test dataset of task `self.task_id`.
110        """
111
112        return MNIST(
113            root=self.root,
114            train=False,
115            transform=self.test_transforms(to_tensor=True),
116            download=False,
117        )