Added new files

utilities/dataset.py

@@ -1,9 +1,7 @@
+import numpy as np
 import pytorch_lightning as pl
-from torch.utils.data import DataLoader
-from torchvision.datasets import CIFAR10
-from torchvision import transforms
 import torch
-import numpy as np
+from torchvision import datasets
 
 
 class CIFAR10(torch.utils.data.Dataset):
@@ -29,33 +27,66 @@ class CIFAR10(torch.utils.data.Dataset):
 
         return (image, label)
 
+
 class CIFAR10DataModule(pl.LightningDataModule):
-    def __init__(self, train_set_transforms,test_set_transforms, data_dir: str = "./data",batch_size: int = 64, num_workers: int = 4):
+    def __init__(
+        self,
+        train_transforms,
+        val_transforms,
+        shuffle=True,
+        data_dir="../data",
+        batch_size=64,
+        num_workers=-1,
+        pin_memory=True,
+    ):
         super().__init__()
+        self.shuffle = shuffle
         self.data_dir = data_dir
         self.batch_size = batch_size
         self.num_workers = num_workers
-        self.train_set_transforms =train_set_transforms
-        self.test_set_transforms = test_set_transforms
+        self.pin_memory = pin_memory
+        self.train_transforms = train_transforms
+        self.val_transforms = val_transforms
+        self.train_data = None
+        self.val_data = None
 
     def prepare_data(self):
-        # Download the CIFAR10 dataset
-        CIFAR10(self.data_dir, train=True, download=True)
-        CIFAR10(self.data_dir, train=False, download=True)
-
-    def setup(self, stage: str = None):
-        # Load the dataset
-        if stage == "fit" or stage is None:
-            self.cifar10_train = CIFAR10(self.data_dir, train=True, transform=self.train_set_transforms)
-            self.cifar10_val = CIFAR10(self.data_dir, train=False, transform=self.train_set_transforms)
-        if stage == "test" or stage is None:
-            self.cifar10_test = CIFAR10(self.data_dir, train=False, transform=self.test_set_transforms)
+        datasets.CIFAR10(self.data_dir, train=True, download=True)
+        datasets.CIFAR10(self.data_dir, train=False, download=True)
+
+    def setup(self, stage):
+        self.train_data = CIFAR10(
+            datasets.CIFAR10(root=self.data_dir, train=True, download=False),
+            transform=self.train_transforms,
+        )
+        self.val_data = CIFAR10(
+            datasets.CIFAR10(root=self.data_dir, train=False, download=False),
+            transform=self.val_transforms,
+        )
 
     def train_dataloader(self):
-        return DataLoader(self.cifar10_train, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=True)
+        return torch.utils.data.DataLoader(
+            self.train_data,
+            batch_size=self.batch_size,
+            shuffle=self.shuffle,
+            num_workers=self.num_workers,
+            pin_memory=self.pin_memory,
+        )
 
     def val_dataloader(self):
-        return DataLoader(self.cifar10_val, batch_size=self.batch_size, num_workers=self.num_workers)
+        return torch.utils.data.DataLoader(
+            self.val_data,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=self.pin_memory,
+        )
 
     def test_dataloader(self):
-        return DataLoader(self.cifar10_test, batch_size=self.batch_size, num_workers=self.num_workers)
+        return torch.utils.data.DataLoader(
+            self.val_data,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=self.pin_memory,
+        )