clarena.stl_datasets.renderedsst2

  1r"""
  2The submodule in `stl_datasets` for Rendered SST2 dataset.
  3"""
  4
  5__all__ = ["RenderedSST2"]
  6
  7import logging
  8from typing import Callable
  9
 10import torch
 11from torch.utils.data import Dataset, random_split
 12from torchvision.datasets import RenderedSST2 as RenderedSST2Raw
 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 RenderedSST2(STLDatasetFromRaw):
 22    r"""Rendered SST2 dataset. The [Rendered SST2 dataset](https://github.com/openai/CLIP/blob/main/data/rendered-sst2.md) is a collection of optical character recognition images. It consists of 9,613 images in 2 classes (positive and negative sentiment), each 448x448 color image."""
 23
 24    original_dataset_python_class: type[Dataset] = RenderedSST2Raw
 25    r"""The original dataset class."""
 26
 27    def __init__(
 28        self,
 29        root: str,
 30        batch_size: int = 1,
 31        num_workers: int = 0,
 32        custom_transforms: Callable | transforms.Compose | None = None,
 33        repeat_channels: int | None = None,
 34        to_tensor: bool = True,
 35        resize: tuple[int, int] | None = None,
 36    ) -> None:
 37        r"""
 38        **Args:**
 39        - **root** (`str`): the root directory where the original Rendered SST2 data 'RenderedSST2/' live.
 40        - **batch_size** (`int`): The batch size in train, val, test dataloader.
 41        - **num_workers** (`int`): the number of workers for dataloaders.
 42        - **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.
 43        - **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.
 44        - **to_tensor** (`bool`): whether to include `ToTensor()` transform. Default is True.
 45        - **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.
 46        """
 47        super().__init__(
 48            root=root,
 49            batch_size=batch_size,
 50            num_workers=num_workers,
 51            custom_transforms=custom_transforms,
 52            repeat_channels=repeat_channels,
 53            to_tensor=to_tensor,
 54            resize=resize,
 55        )
 56
 57    def prepare_data(self) -> None:
 58        r"""Download the original Rendered SST2 dataset if haven't."""
 59
 60        RenderedSST2Raw(root=self.root, split="train", download=True)
 61        RenderedSST2Raw(root=self.root, split="val", download=True)
 62        RenderedSST2Raw(root=self.root, split="test", download=True)
 63
 64        pylogger.debug(
 65            "The original Rendered SST2 dataset has been downloaded to %s.",
 66            self.root,
 67        )
 68
 69    def train_and_val_dataset(self) -> tuple[Dataset, Dataset]:
 70        """Get the training and validation dataset.
 71
 72        **Returns:**
 73        - **train_and_val_dataset** (`tuple[Dataset, Dataset]`): the train and validation dataset.
 74        """
 75        dataset_train = RenderedSST2Raw(
 76            root=self.root,
 77            split="train",
 78            transform=self.train_and_val_transforms(),
 79            target_transform=self.target_transform(),
 80            download=False,
 81        )
 82
 83        dataset_val = RenderedSST2Raw(
 84            root=self.root,
 85            split="val",
 86            transform=self.train_and_val_transforms(),
 87            target_transform=self.target_transform(),
 88            download=False,
 89        )
 90
 91        return dataset_train, dataset_val
 92
 93    def test_dataset(self) -> Dataset:
 94        r"""Get the test dataset.
 95
 96        **Returns:**
 97        - **test_dataset** (`Dataset`): the test dataset.
 98        """
 99        dataset_test = RenderedSST2Raw(
100            root=self.root,
101            split="test",
102            transform=self.test_transforms(),
103            target_transform=self.target_transform(),
104            download=False,
105        )
106
107        return dataset_test