my_project.dataset

 1"""
 2Defines the PyTorch Lightning DataModule for the Fashion-MNIST dataset.
 3
 4This module is responsible for:
 5- Downloading the Fashion-MNIST dataset.
 6- Applying necessary transformations (e.g., normalization).
 7- Splitting the data into training, validation, and test sets.
 8- Creating and providing DataLoaders for each set.
 9"""
10
11import pytorch_lightning as pl
12from torch.utils.data import DataLoader, random_split
13from torchvision import datasets, transforms
14
15
16class FashionMNISTDataModule(pl.LightningDataModule):
17    """
18    PyTorch Lightning DataModule for the Fashion-MNIST dataset.
19    It handles the downloading, splitting, and loading of the data.
20    """
21
22    def __init__(self, data_dir: str = "data/", batch_size: int = 128, val_split: float = 0.2, num_workers: int = 4):
23        """
24        Args:
25            data_dir (str): Directory where the data will be downloaded/stored.
26            batch_size (int): The batch size for the data loaders.
27            val_split (float): The fraction of the training data to use for validation.
28            num_workers (int): Number of subprocesses to use for data loading.
29        """
30        super().__init__()
31        self.data_dir = data_dir
32        self.batch_size = batch_size
33        self.num_workers = num_workers # Store num_workers
34        self.val_split = val_split
35        self.transform = transforms.Compose(
36            [
37                transforms.ToTensor(),
38                transforms.Normalize((0.5,), (0.5,)),  # Normalize to [-1, 1]
39            ]
40        )
41        self.train_dataset = None
42        self.val_dataset = None
43        self.test_dataset = None
44        self.train_val_dataset = None # For data exploration tab
45
46    def prepare_data(self):
47        """
48        Downloads the Fashion-MNIST dataset if it's not already present.
49        This method is called only on a single GPU/process.
50        """
51        datasets.FashionMNIST(self.data_dir, train=True, download=True)
52        datasets.FashionMNIST(self.data_dir, train=False, download=True)
53
54    def setup(self, stage: str = None):
55        """
56        Assigns train/val/test datasets for dataloaders.
57        This method is called on every GPU.
58        """
59        # Assign train/val datasets for use in dataloaders
60        if stage == "fit" or stage is None:
61            self.train_val_dataset = datasets.FashionMNIST(self.data_dir, train=True, transform=self.transform)
62            n_samples = len(self.train_val_dataset)
63            n_val = int(self.val_split * n_samples)
64            n_train = n_samples - n_val
65            self.train_dataset, self.val_dataset = random_split(self.train_val_dataset, [n_train, n_val])
66
67        # Assign test dataset for use in dataloader(s)
68        if stage == "test" or stage is None:
69            self.test_dataset = datasets.FashionMNIST(self.data_dir, train=False, transform=self.transform)
70
71    def train_dataloader(self):
72        return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=self.num_workers)
73
74    def val_dataloader(self):
75        return DataLoader(self.val_dataset, batch_size=self.batch_size, num_workers=self.num_workers)
76
77    def test_dataloader(self):
78        return DataLoader(self.test_dataset, batch_size=self.batch_size, num_workers=self.num_workers)