feat: add tiny datasets for lightweight experiments

pyproject.toml

@@ -35,6 +35,7 @@ possibly-missing-attribute = "ignore" # mypy is more permissive with attribute a
 possibly-missing-import = "ignore" # mypy is more permissive with imports
 no-matching-overload = "ignore" # mypy is more permissive with overloads
 unresolved-reference = "ignore" # mypy is more permissive with references
+missing-argument = "ignore"
 
 [tool.coverage.run]
 source = ["src/pruna"]

src/pruna/data/__init__.py

@@ -23,6 +23,7 @@
     setup_cifar10_dataset,
     setup_imagenet_dataset,
     setup_mnist_dataset,
+    setup_tiny_cifar10_dataset,
 )
 from pruna.data.datasets.prompt import (
     setup_drawbench_dataset,
@@ -77,6 +78,7 @@
         "image_classification_collate",
         {"img_size": 32},
     ),
+    "TinyCIFAR10": (setup_tiny_cifar10_dataset, "image_classification_collate", {"img_size": 32}),
     "DrawBench": (setup_drawbench_dataset, "prompt_collate", {}),
     "PartiPrompts": (setup_parti_prompts_dataset, "prompt_collate", {}),
     "GenAIBench": (setup_genai_bench_dataset, "prompt_collate", {}),

src/pruna/data/datasets/image.py

@@ -66,6 +66,10 @@ def setup_cifar10_dataset(seed: int) -> Tuple[Dataset, Dataset, Dataset]:
     """
     Setup the CIFAR-10 dataset.
 
+    The original CIFAR-10 dataset from uoft-cs/cifar10 has an 'img' column,
+    but this function renames it to 'image' to ensure compatibility with
+    the image_classification_collate function which expects an 'image' column.
+
     License: unspecified
 
     Parameters
@@ -76,8 +80,48 @@ def setup_cifar10_dataset(seed: int) -> Tuple[Dataset, Dataset, Dataset]:
     Returns
     -------
     Tuple[Dataset, Dataset, Dataset]
-        The CIFAR-10 dataset.
+        The CIFAR-10 dataset with columns: 'image' (PIL Image) and 'label' (int).
     """
-    train_ds, test_ds = load_dataset("uoft-cs/cifar10", split=["train", "test"])  # type: ignore[misc]
+    train_ds, test_ds = load_dataset("uoft-cs/cifar10", split=["train", "test"])
+
+    # Rename 'img' column to 'image' to match collate function expectations
+    # This ensures compatibility with image_classification_collate function
+    train_ds = train_ds.rename_column("img", "image")
+    test_ds = test_ds.rename_column("img", "image")
+
     train_ds, val_ds = split_train_into_train_val(train_ds, seed)
     return train_ds, val_ds, test_ds  # type: ignore[return-value]
+
+
+def setup_tiny_cifar10_dataset(seed: int) -> Tuple[Dataset, Dataset, Dataset]:
+    """
+    Setup the Tiny CIFAR-10 dataset (< 1,000 samples).
+
+    The original CIFAR-10 dataset from uoft-cs/cifar10 has an 'img' column,
+    but this function renames it to 'image' to ensure compatibility with
+    the image_classification_collate function which expects an 'image' column.
+
+    License: unspecified
+
+    Parameters
+    ----------
+    seed : int
+        The seed to use.
+
+    Returns
+    -------
+    Tuple[Dataset, Dataset, Dataset]
+        The Tiny CIFAR-10 dataset with columns: 'image' (PIL Image) and 'label' (int).
+        Contains approximately 600 training samples, split validation, and 200 test samples.
+    """
+    train_ds, test_ds = load_dataset("uoft-cs/cifar10", split=["train", "test"])
+
+    # Rename 'img' column to 'image' to match collate function expectations
+    # This ensures compatibility with image_classification_collate function
+    train_ds = train_ds.rename_column("img", "image")
+    test_ds = test_ds.rename_column("img", "image")
+
+    tiny_train = train_ds.select(range(600))
+    tiny_test = test_ds.select(range(200))
+    train_ds, val_ds = split_train_into_train_val(tiny_train, seed)
+    return train_ds, val_ds, tiny_test

tests/data/test_datamodule.py

@@ -28,6 +28,7 @@ def iterate_dataloaders(datamodule: PrunaDataModule) -> None:
         pytest.param("LibriSpeech", dict(), marks=pytest.mark.slow),
         pytest.param("AIPodcast", dict(), marks=pytest.mark.slow),
         ("ImageNet", dict(img_size=512)),
+        ("TinyCIFAR10", dict(img_size=32)),
         pytest.param("MNIST", dict(img_size=512), marks=pytest.mark.slow),
         ("WikiText", dict(tokenizer=bert_tokenizer)),
         pytest.param("TinyWikiText", dict(tokenizer=bert_tokenizer), marks=pytest.mark.slow),