Add CIFAR10 dataset and input/output shape getters.

Author: mitkonikov

src/fflib/utils/data/cifar10.py

@@ -0,0 +1,137 @@
+import torch
+
+from torch.nn.functional import one_hot
+from torch.utils.data import DataLoader, random_split
+from torchvision.datasets import CIFAR10  # type: ignore
+from torchvision.transforms import Compose, ToTensor, Lambda  # type: ignore
+
+from fflib.utils.data import FFDataProcessor
+from fflib.interfaces.iff import IFF
+
+from enum import Enum
+from typing import Tuple, Dict, Callable, Any
+
+
+class NegativeGenerator(Enum):
+    INVERSE = 1
+    RANDOM = 2
+    HIGHEST_INCORRECT = 3
+
+
+class FFCIFAR10(FFDataProcessor):
+    def __init__(
+        self,
+        batch_size: int,
+        validation_split: float | None,
+        download: bool = True,
+        path: str = "./data",
+        image_transform: Callable[..., Any] = Compose([ToTensor(), Lambda(torch.flatten)]),
+        train_kwargs: Dict[str, Any] = {},
+        test_kwargs: Dict[str, Any] = {},
+        negative_generator: NegativeGenerator = NegativeGenerator.INVERSE,
+        use: float = 1.0,
+    ):
+
+        assert isinstance(batch_size, int)
+        assert batch_size > 0
+        self.batch_size = batch_size
+        if "batch_size" not in train_kwargs:
+            train_kwargs["batch_size"] = self.batch_size
+        if "batch_size" not in test_kwargs:
+            test_kwargs["batch_size"] = self.batch_size
+
+        train_kwargs["shuffle"] = True
+
+        assert use >= 0.0 and use <= 1.0
+
+        self.validation_split = validation_split
+        self.download = download
+        self.path = path
+        self.image_transform = image_transform
+        self.train_kwargs = train_kwargs
+        self.test_kwargs = test_kwargs
+        self.negative_generator = negative_generator
+        self.use = use
+
+        self.train_dataset = CIFAR10(
+            self.path, train=True, download=self.download, transform=self.image_transform
+        )
+        self.test_dataset = CIFAR10(
+            self.path, train=False, download=self.download, transform=self.image_transform
+        )
+        self.test_loader = DataLoader(self.test_dataset, **self.test_kwargs)
+
+        dataset_size = len(self.train_dataset)
+        used_dataset_size = int(dataset_size * self.use)
+        not_used_dataset_size = dataset_size - used_dataset_size
+
+        # In case a validation split is given
+        if self.validation_split:
+            # Determine the sizes of training and validation sets
+            val_size = int(self.validation_split * used_dataset_size)
+            train_size = used_dataset_size - val_size
+
+            # Split dataset into train and validation sets
+            train_dataset, val_dataset, _ = random_split(
+                self.train_dataset, [train_size, val_size, not_used_dataset_size]
+            )
+
+            # Create data loaders for train and validation
+            self.train_loader = DataLoader(train_dataset, **self.train_kwargs)
+            self.val_loader = DataLoader(val_dataset, **self.test_kwargs)
+
+            assert len(self.train_loader) + len(self.val_loader) <= used_dataset_size
+            return
+
+        train_dataset, _ = random_split(
+            self.train_dataset, [used_dataset_size, not_used_dataset_size]
+        )
+        self.train_loader = DataLoader(train_dataset, **self.train_kwargs)
+
+    def get_input_shape(self) -> torch.Size:
+        return torch.Size((32 * 32 * 3,))
+
+    def get_output_shape(self) -> torch.Size:
+        return torch.Size((10,))
+
+    def get_train_loader(self) -> DataLoader[Any]:
+        return self.train_loader
+
+    def get_val_loader(self) -> DataLoader[Any]:
+        return self.val_loader
+
+    def get_test_loader(self) -> DataLoader[Any]:
+        return self.test_loader
+
+    def get_all_loaders(self) -> Dict[str, DataLoader[Any]]:
+        return {
+            "train": self.get_train_loader(),
+            "val": self.get_val_loader(),
+            "test": self.get_test_loader(),
+        }
+
+    def encode_output(self, y: torch.Tensor) -> torch.Tensor:
+        return one_hot(y, num_classes=10).float()
+
+    def combine_to_input(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        return torch.cat((x, y), 1)
+
+    def generate_negative(
+        self,
+        x: torch.Tensor,
+        y: torch.Tensor,
+        net: IFF,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+
+        if self.negative_generator == NegativeGenerator.HIGHEST_INCORRECT:
+            raise NotImplementedError()
+
+        if self.negative_generator == NegativeGenerator.INVERSE:
+            y_hot = 1 - one_hot(y, num_classes=10).float()
+            return x, y_hot
+
+        rnd = torch.rand((x.shape[0], 10), device=x.device)
+        rnd[torch.arange(x.shape[0]), y] = 0
+        y_new = rnd.argmax(1)
+        y_hot = one_hot(y_new, num_classes=10).float()
+        return x, y_hot

src/fflib/utils/data/dataprocessor.py

@@ -15,6 +15,17 @@ def check_splits(splits: Tuple[float, float, float] | Tuple[float, float]) -> No
         assert all(0 <= s <= 1 for s in splits)
         assert len(splits) in [2, 3]
 
+    @abstractmethod
+    def get_input_shape(self) -> torch.Size:
+        pass
+
+    @abstractmethod
+    def get_output_shape(self) -> torch.Size:
+        pass
+
+    def numel(self) -> int:
+        return self.get_input_shape().numel() + self.get_output_shape().numel()
+
     @abstractmethod
     def get_train_loader(self) -> DataLoader[Any]:
         pass

src/fflib/utils/data/datasets.py

@@ -1,8 +1,9 @@
 from fflib.utils.data.mnist import FFMNIST
 from fflib.utils.data.fashion_mnist import FFFashionMNIST
 from fflib.utils.data.xor import FFXOR
+from fflib.utils.data.cifar10 import FFCIFAR10
 
-from typing import Dict, Any
+from typing import Any
 
 
 def CreateDatasetFromName(
@@ -11,7 +12,7 @@ def CreateDatasetFromName(
     validation_split: float,
     use: float = 1.0,
     **kwargs: Any,
-) -> FFMNIST | FFFashionMNIST | FFXOR | None:
+) -> FFMNIST | FFFashionMNIST | FFXOR | FFCIFAR10 | None:
     """Create a Dataset object via a name of the dataset.
     This is used to dynamically set the dataset from CLI arguments.
 
@@ -41,6 +42,13 @@ def CreateDatasetFromName(
             use=use,
             **kwargs,
         )
+    elif name == "cifar10":
+        return FFCIFAR10(
+            batch_size,
+            validation_split,
+            use=use,
+            **kwargs,
+        )
     elif name == "xor":
         return FFXOR(
             batch_size,

src/fflib/utils/data/fashion_mnist.py

@@ -88,6 +88,12 @@ def __init__(
         )
         self.train_loader = DataLoader(train_dataset, **self.train_kwargs)
 
+    def get_input_shape(self) -> torch.Size:
+        return torch.Size((28 * 28,))
+
+    def get_output_shape(self) -> torch.Size:
+        return torch.Size((10,))
+
     def get_train_loader(self) -> DataLoader[Any]:
         return self.train_loader
 

src/fflib/utils/data/mnist.py

@@ -88,6 +88,12 @@ def __init__(
         )
         self.train_loader = DataLoader(train_dataset, **self.train_kwargs)
 
+    def get_input_shape(self) -> torch.Size:
+        return torch.Size((28 * 28,))
+
+    def get_output_shape(self) -> torch.Size:
+        return torch.Size((10,))
+
     def get_train_loader(self) -> DataLoader[Any]:
         return self.train_loader
 

src/fflib/utils/data/xor.py

@@ -48,6 +48,12 @@ def __init__(
         self.test_loader = DataLoader(XORDataset(size), **self.test_kwargs)
         self.val_loader = DataLoader(XORDataset(size), **self.test_kwargs)
 
+    def get_input_shape(self) -> torch.Size:
+        return torch.Size((2,))
+
+    def get_output_shape(self) -> torch.Size:
+        return torch.Size((2,))
+
     def get_train_loader(self) -> DataLoader[Any]:
         return self.train_loader