clarena.stl_datasets.SEMEION

  1r"""
  2The submodule in `stl_datasets` for SEMEION dataset.
  3"""
  4
  5__all__ = ["SEMEION"]
  6
  7import logging
  8from typing import Callable
  9
 10import torch
 11from torch.utils.data import Dataset, random_split
 12from torchvision.datasets import SEMEION as SEMEIONRaw
 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 SEMEION(STLDatasetFromRaw):
 22    r"""SEMEION dataset. The [SEMEION dataset](https://archive.ics.uci.edu/dataset/178/semeion+handwritten+digit) is a collection of handwritten digits. It consists of 1,593 handwritten digit images (10 classes), each 16x16 grayscale image."""
 23
 24    original_dataset_python_class: type[Dataset] = SEMEIONRaw
 25    r"""The original dataset class."""
 26
 27    def __init__(
 28        self,
 29        root: str,
 30        test_percentage: float,
 31        validation_percentage: float,
 32        batch_size: int = 1,
 33        num_workers: int = 0,
 34        custom_transforms: Callable | transforms.Compose | None = None,
 35        repeat_channels: int | None = None,
 36        to_tensor: bool = True,
 37        resize: tuple[int, int] | None = None,
 38    ) -> None:
 39        r"""
 40        **Args:**
 41        - **root** (`str`): the root directory where the original SEMEION data 'SEMEION/' live.
 42        - **test_percentage** (`float`): the percentage to randomly split some data into test data.
 43        - **validation_percentage** (`float`): the percentage to randomly split some training data into validation data.
 44        - **batch_size** (`int`): The batch size in train, val, test dataloader.
 45        - **num_workers** (`int`): the number of workers for dataloaders.
 46        - **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.
 47        - **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.
 48        - **to_tensor** (`bool`): whether to include `ToTensor()` transform. Default is True.
 49        - **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.
 50        """
 51        super().__init__(
 52            root=root,
 53            batch_size=batch_size,
 54            num_workers=num_workers,
 55            custom_transforms=custom_transforms,
 56            repeat_channels=repeat_channels,
 57            to_tensor=to_tensor,
 58            resize=resize,
 59        )
 60
 61        self.test_percentage: float = test_percentage
 62        r"""The percentage to randomly split some data into test data."""
 63        self.validation_percentage: float = validation_percentage
 64        r"""The percentage to randomly split some training data into validation data."""
 65
 66    def prepare_data(self) -> None:
 67        r"""Download the original SEMEION dataset if haven't."""
 68
 69        SEMEIONRaw(root=self.root, download=True)
 70
 71        pylogger.debug(
 72            "The original SEMEION dataset has been downloaded to %s.", self.root
 73        )
 74
 75    def train_and_val_dataset(self) -> tuple[Dataset, Dataset]:
 76        """Get the training and validation dataset.
 77
 78        **Returns:**
 79        - **train_and_val_dataset** (`tuple[Dataset, Dataset]`): the train and validation dataset.
 80        """
 81        dataset_all = SEMEIONRaw(
 82            root=self.root,
 83            transform=self.train_and_val_transforms(),
 84            target_transform=self.target_transform(),
 85            download=False,
 86        )
 87
 88        dataset_train_and_val, _ = random_split(
 89            dataset_all,
 90            lengths=[
 91                1 - self.test_percentage,
 92                self.test_percentage,
 93            ],
 94            generator=torch.Generator().manual_seed(
 95                42
 96            ),  # 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
 97        )
 98
 99        return random_split(
100            dataset_train_and_val,
101            lengths=[1 - self.validation_percentage, self.validation_percentage],
102            generator=torch.Generator().manual_seed(
103                42
104            ),  # 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
105        )
106
107    def test_dataset(self) -> Dataset:
108        r"""Get the test dataset.
109
110        **Returns:**
111        - **test_dataset** (`Dataset`): the test dataset.
112        """
113        dataset_all = SEMEIONRaw(
114            root=self.root,
115            transform=self.train_and_val_transforms(),
116            target_transform=self.target_transform(),
117            download=False,
118        )
119
120        _, dataset_test = random_split(
121            dataset_all,
122            lengths=[1 - self.test_percentage, self.test_percentage],
123            generator=torch.Generator().manual_seed(42),
124        )
125
126        return dataset_test