clarena.stl_datasets.kmnist

  1r"""
  2The submodule in `stl_datasets` for KMNIST dataset.
  3"""
  4
  5__all__ = ["KMNIST"]
  6
  7import logging
  8from typing import Callable
  9
 10import torch
 11from torch.utils.data import Dataset, random_split
 12from torchvision.datasets import KMNIST as KMNISTRaw
 13from torchvision.transforms import transforms
 14
 15from clarena.stl_datasets.base import STLDatasetFromRaw
 16
 17# always get logger for built-in logging in each module
 18pylogger = logging.getLogger(__name__)
 19
 20
 21class KMNIST(STLDatasetFromRaw):
 22    r"""Kuzushiji-MNIST dataset. The [Kuzushiji-MNIST dataset](https://github.com/rois-codh/kmnist) is a collection of Japanese Kuzushiji character images. It consists of 60,000 training and 10,000 test images of Japanese Kuzushiji images (10 classes), each 28x28 grayscale image (similar to MNIST)."""
 23
 24    original_dataset_python_class: type[Dataset] = KMNISTRaw
 25    r"""The original dataset class."""
 26
 27    def __init__(
 28        self,
 29        root: str,
 30        validation_percentage: float,
 31        batch_size: int = 1,
 32        num_workers: int = 0,
 33        custom_transforms: Callable | transforms.Compose | None = None,
 34        repeat_channels: int | None = None,
 35        to_tensor: bool = True,
 36        resize: tuple[int, int] | None = None,
 37    ) -> None:
 38        r"""
 39        **Args:**
 40        - **root** (`str`): the root directory where the original Kuzushiji-MNIST data 'KMNIST/' live.
 41        - **validation_percentage** (`float`): the percentage to randomly split some training data into validation data.
 42        - **batch_size** (`int`): The batch size in train, val, test dataloader.
 43        - **num_workers** (`int`): the number of workers for dataloaders.
 44        - **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. `ToTensor()`, normalize and so on are not included.
 45        - **repeat_channels** (`int` | `None`): the number of channels to repeat. Default is None, which means no repeat. If not None, it should be an integer.
 46        - **to_tensor** (`bool`): whether to include `ToTensor()` transform. Default is True.
 47        - **resize** (`tuple[int, int]` | `None` or list of them): the size to resize the images to. Default is None, which means no resize. If not None, it should be a tuple of two integers.
 48        """
 49        super().__init__(
 50            root=root,
 51            batch_size=batch_size,
 52            num_workers=num_workers,
 53            custom_transforms=custom_transforms,
 54            repeat_channels=repeat_channels,
 55            to_tensor=to_tensor,
 56            resize=resize,
 57        )
 58
 59        self.validation_percentage: float = validation_percentage
 60        r"""The percentage to randomly split some training data into validation data."""
 61
 62    def prepare_data(self) -> None:
 63        r"""Download the original Kuzushiji-MNIST dataset if haven't."""
 64
 65        KMNISTRaw(root=self.root, train=True, download=True)
 66        KMNISTRaw(root=self.root, train=False, download=True)
 67
 68        pylogger.debug(
 69            "The original Kuzushiji-MNIST dataset has been downloaded to %s.",
 70            self.root,
 71        )
 72
 73    def train_and_val_dataset(self) -> tuple[Dataset, Dataset]:
 74        """Get the training and validation dataset.
 75
 76        **Returns:**
 77        - **train_and_val_dataset** (`tuple[Dataset, Dataset]`): the train and validation dataset.
 78        """
 79        dataset_train_and_val = KMNISTRaw(
 80            root=self.root,
 81            train=True,
 82            transform=self.train_and_val_transforms(),
 83            target_transform=self.target_transform(),
 84            download=False,
 85        )
 86
 87        return random_split(
 88            dataset_train_and_val,
 89            lengths=[1 - self.validation_percentage, self.validation_percentage],
 90            generator=torch.Generator().manual_seed(
 91                42
 92            ),  # 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
 93        )
 94
 95    def test_dataset(self) -> Dataset:
 96        r"""Get the test dataset.
 97
 98        **Returns:**
 99        - **test_dataset** (`Dataset`): the test dataset.
100        """
101        dataset_test = KMNISTRaw(
102            root=self.root,
103            train=False,
104            transform=self.test_transforms(),
105            target_transform=self.target_transform(),
106            download=False,
107        )
108
109        return dataset_test