clarena.cl_datasets.permuted_mnist

  1"""
  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    """Permuted MNIST dataset."""
 23
 24    num_class: 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 deviation 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 = 10,
 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        """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        super().__init__(
 66            root=root,
 67            num_tasks=num_tasks,
 68            validation_percentage=validation_percentage,
 69            batch_size=batch_size,
 70            num_workers=num_workers,
 71            custom_transforms=custom_transforms,
 72            custom_target_transforms=custom_target_transforms,
 73            permutation_mode=permutation_mode,
 74            permutation_seeds=permutation_seeds,
 75        )
 76
 77    def prepare_data(self) -> None:
 78        """Download the original MNIST dataset if haven't."""
 79        # just download
 80        MNIST(root=self.root, train=True, download=True)
 81        MNIST(root=self.root, train=False, download=True)
 82
 83        pylogger.debug(
 84            "The original MNIST dataset has been downloaded to %s.", self.root
 85        )
 86
 87    def train_and_val_dataset(self) -> tuple[Dataset, Dataset]:
 88        """Get the training and validation dataset of task `self.task_id`.
 89
 90        **Returns:**
 91        - The train and validation dataset of task `self.task_id`."""
 92        dataset_train_and_val = MNIST(
 93            root=self.root,
 94            train=True,
 95            transform=self.train_and_val_transforms(to_tensor=True),
 96            target_transform=self.target_transforms(),
 97            download=False,
 98        )
 99        return random_split(
100            dataset_train_and_val,
101            lengths=[1 - self.validation_percentage, self.validation_percentage],
102        )
103
104    def test_dataset(self) -> Dataset:
105        """Get the test dataset of task `self.task_id`.
106
107        **Returns:**
108        - The test dataset of task `self.task_id`.
109        """
110        return MNIST(
111            root=self.root,
112            train=False,
113            transform=self.test_transforms(to_tensor=True),
114            target_transform=self.target_transforms(),
115            download=False,
116        )