feat: add stratifying by sample size for image classification datasets

Author: begumcig

src/pruna/data/__init__.py

@@ -78,9 +78,24 @@
         "image_classification_collate",
         {"img_size": 32},
     ),
-    "TinyCIFAR10": (partial(setup_cifar10_dataset, fraction=0.1), "image_classification_collate", {"img_size": 32}),
-    "TinyMNIST": (partial(setup_mnist_dataset, fraction=0.1), "image_classification_collate", {"img_size": 28}),
-    "TinyImageNet": (partial(setup_imagenet_dataset, fraction=0.1), "image_classification_collate", {"img_size": 224}),
+    # our full CIFAR10 has 50k train and 10k test
+    "TinyCIFAR10": (
+        partial(setup_cifar10_dataset, train_sample_size=800, test_sample_size=100),
+        "image_classification_collate",
+        {"img_size": 32},
+    ),
+    #  our full MNIST has 60k train and 10k test
+    "TinyMNIST": (
+        partial(setup_mnist_dataset, train_sample_size=800, test_sample_size=100),
+        "image_classification_collate",
+        {"img_size": 28},
+    ),
+    # our full ImageNet has 100k train and 10k val
+    "TinyImageNet": (
+        partial(setup_imagenet_dataset, train_sample_size=1000, test_sample_size=100),
+        "image_classification_collate",
+        {"img_size": 224},
+    ),
     "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

@@ -12,15 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 from typing import Tuple
 
 from datasets import load_dataset
 from torch.utils.data import Dataset
 
-from pruna.data.utils import split_train_into_train_val, split_val_into_val_test, stratify_dataset
+from pruna.data.utils import (
+    define_sample_size_for_dataset,
+    split_train_into_train_val,
+    split_val_into_val_test,
+    stratify_dataset,
+)
 
 
-def setup_mnist_dataset(seed: int, fraction: float = 1.0) -> Tuple[Dataset, Dataset, Dataset]:
+def setup_mnist_dataset(
+    seed: int, fraction: float = 1.0, train_sample_size: int | None = None, test_sample_size: int | None = None
+) -> Tuple[Dataset, Dataset, Dataset]:
     """
     Setup the MNIST dataset.
 
@@ -30,9 +39,12 @@ def setup_mnist_dataset(seed: int, fraction: float = 1.0) -> Tuple[Dataset, Data
     ----------
     seed : int
         The seed to use.
-
     fraction : float
         The fraction of the dataset to use.
+    train_sample_size : int | None
+        The sample size to use for the train dataset.
+    test_sample_size : int | None
+        The sample size to use for the test dataset.
 
     Returns
     -------
@@ -41,16 +53,21 @@ def setup_mnist_dataset(seed: int, fraction: float = 1.0) -> Tuple[Dataset, Data
     """
     train_ds, test_ds = load_dataset("ylecun/mnist", split=["train", "test"])  # type: ignore[misc]
 
-    train_ds = stratify_dataset(train_ds, "label", fraction, seed)
-    test_ds = stratify_dataset(test_ds, "label", fraction, seed)
+    train_sample_size = define_sample_size_for_dataset(train_ds, fraction, train_sample_size)
+    test_sample_size = define_sample_size_for_dataset(test_ds, fraction, test_sample_size)
+
+    train_ds = stratify_dataset(train_ds, train_sample_size, seed)
+    test_ds = stratify_dataset(test_ds, test_sample_size, seed)
 
     train_ds, val_ds = split_train_into_train_val(train_ds, seed)
     val_ds, test_ds = split_val_into_val_test(val_ds, seed)
 
     return train_ds, val_ds, test_ds  # type: ignore[return-value]
 
 
-def setup_imagenet_dataset(seed: int, fraction: float = 1.0) -> Tuple[Dataset, Dataset, Dataset]:
+def setup_imagenet_dataset(
+    seed: int, fraction: float = 1.0, train_sample_size: int | None = None, test_sample_size: int | None = None
+) -> Tuple[Dataset, Dataset, Dataset]:
     """
     Setup the ImageNet dataset.
 
@@ -60,23 +77,30 @@ def setup_imagenet_dataset(seed: int, fraction: float = 1.0) -> Tuple[Dataset, D
     ----------
     seed : int
         The seed to use.
-
     fraction : float
         The fraction of the dataset to use.
+    train_sample_size : int | None
+        The sample size to use for the train dataset.
+    test_sample_size : int | None
+        The sample size to use for the test dataset.
 
     Returns
     -------
     Tuple[Dataset, Dataset, Dataset]
         The ImageNet dataset.
     """
     train_ds, val = load_dataset("zh-plus/tiny-imagenet", split=["train", "valid"])  # type: ignore[misc]
-    train_ds = stratify_dataset(train_ds, "label", fraction, seed)
-    val = stratify_dataset(val, "label", fraction, seed)
+    train_sample_size = define_sample_size_for_dataset(train_ds, fraction, train_sample_size)
+    train_ds = stratify_dataset(train_ds, train_sample_size, seed)
     val_ds, test_ds = split_val_into_val_test(val, seed)
+    test_sample_size = define_sample_size_for_dataset(test_ds, fraction, test_sample_size)
+    test_ds = stratify_dataset(test_ds, test_sample_size, seed)
     return train_ds, val_ds, test_ds  # type: ignore[return-value]
 
 
-def setup_cifar10_dataset(seed: int, fraction: float = 1.0) -> Tuple[Dataset, Dataset, Dataset]:
+def setup_cifar10_dataset(
+    seed: int, fraction: float = 1.0, train_sample_size: int | None = None, test_sample_size: int | None = None
+) -> Tuple[Dataset, Dataset, Dataset]:
     """
     Setup the CIFAR-10 dataset.
 
@@ -90,9 +114,12 @@ def setup_cifar10_dataset(seed: int, fraction: float = 1.0) -> Tuple[Dataset, Da
     ----------
     seed : int
         The seed to use.
-
     fraction : float
         The fraction of the dataset to use.
+    train_sample_size : int | None
+        The sample size to use for the train dataset.
+    test_sample_size : int | None
+        The sample size to use for the test dataset.
 
     Returns
     -------
@@ -106,8 +133,11 @@ def setup_cifar10_dataset(seed: int, fraction: float = 1.0) -> Tuple[Dataset, Da
     train_ds = train_ds.rename_column("img", "image")
     test_ds = test_ds.rename_column("img", "image")
 
-    train_ds = stratify_dataset(train_ds, "label", fraction, seed)
-    test_ds = stratify_dataset(test_ds, "label", fraction, seed)
+    train_sample_size = define_sample_size_for_dataset(train_ds, fraction, train_sample_size)
+    test_sample_size = define_sample_size_for_dataset(test_ds, fraction, test_sample_size)
+
+    train_ds = stratify_dataset(train_ds, train_sample_size, seed)
+    test_ds = stratify_dataset(test_ds, test_sample_size, seed)
 
     train_ds, val_ds = split_train_into_train_val(train_ds, seed)
     return train_ds, val_ds, test_ds  # type: ignore[return-value]

src/pruna/data/utils.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import random
 from typing import Any, Tuple, Union
 
 import torch
@@ -183,28 +184,58 @@ def recover_text_from_dataloader(dataloader: DataLoader, tokenizer: Any) -> list
     return texts
 
 
-def stratify_dataset(dataset: Dataset, column: str, fraction: float, seed: int) -> Dataset:
+def stratify_dataset(dataset: Dataset, sample_size: int, seed: int = 42) -> Dataset:
     """
-    Stratify the dataset into a fraction of the dataset.
+    Stratify the dataset into a specific size.
 
     Parameters
     ----------
     dataset : Dataset
         The dataset to stratify.
-    column : str
-        The column to stratify by.
-    fraction : float
-        The fraction of the dataset to stratify.
+    sample_size : int
+        The size to stratify.
     seed : int
-        The seed to use for splitting the dataset.
+        The seed to use for sampling the dataset.
 
     Returns
     -------
     Dataset
         The stratified dataset.
     """
-    if fraction < 1.0:
-        split_result = dataset.train_test_split(test_size=1 - fraction, stratify_by_column="label", seed=seed)
-        dataset = split_result["train"]
-
+    dataset_length = len(dataset)
+    if dataset_length < sample_size:
+        pruna_logger.warning(
+            "Dataset length is less than the size to stratify."
+            f"Using the entire dataset. ({dataset_length} < {sample_size})"
+        )
+        return dataset
+
+    indices = list(range(dataset_length))
+    random.Random(seed).shuffle(indices)
+    selected_indices = indices[:sample_size]
+    dataset = dataset.select(selected_indices)
     return dataset
+
+
+def define_sample_size_for_dataset(dataset: Dataset, fraction: float, sample_size: int | None = None) -> int:
+    """
+    Define the sample size for the dataset.
+
+    Parameters
+    ----------
+    dataset: Dataset
+        The dataset to define the sample size for.
+    fraction: float
+        The fraction of the dataset to sample.
+    sample_size: int | None
+        The sample size to use.
+
+    Returns
+    -------
+    int
+        The sample size for the dataset.
+    """
+    if fraction < 1.0 and (sample_size is not None):
+        raise ValueError("Fraction and sample sizes cannot be used together.")
+    sample_size = int(len(dataset) * fraction) if fraction < 1.0 else sample_size or len(dataset)
+    return sample_size