Interaction Nvidia DALI and Squirrel

docs/conf.py

@@ -48,7 +48,9 @@
     "sphinx.ext.todo",
     "autoapi.extension",
     "sphinxcontrib.mermaid",
+    "myst_nb",
 ]
+nb_execution_mode = "off"
 
 # Add any paths that contain templates here, relative to this directory.
 # Add any paths that contain templates here, relative to this directory.cv

docs/examples/create_imagenette.py

@@ -0,0 +1,38 @@
+import typing as t
+
+import numpy as np
+import torchvision.transforms.functional as F
+from create_sq_ds_from_hf import create_sq_ds_from_hf, to_np_dict
+
+
+def to_np_dict_imagenette(
+    item: t.Dict[str, t.Any],
+    output_size: int = 300,
+    img_key: str = "image",
+) -> t.Dict[str, np.ndarray]:
+    """Converts imagenette images to a fixed, square size."""
+
+    item = to_np_dict(item)
+
+    img = F.to_tensor(item[img_key])  # converts from HWC to CHW
+    img = F.resize(img, output_size)
+    img = F.center_crop(img, output_size)
+    img = img.permute((1, 2, 0))  # convert back from CHW to HWC
+
+    # handle grayscale images in imagenette which cause trouble downstream
+    if img.shape[2] == 1:
+        img = img.repeat(1, 1, 3)  # repeat channel dimension 3 times
+
+    item[img_key] = img.numpy()
+
+    return item
+
+
+if __name__ == "__main__":
+    create_sq_ds_from_hf(
+        "frgfm/imagenette",
+        to_np_dict=to_np_dict_imagenette,
+        shard_size=10000,
+        subset="320px",
+        hf_split="train",
+    )

docs/examples/create_sq_ds_from_hf.py

@@ -0,0 +1,54 @@
+import typing as t
+from pathlib import Path
+
+import numpy as np
+
+from squirrel.driver.msgpack import MessagepackDriver
+from squirrel_datasets_core.driver.huggingface import HuggingfaceDriver
+
+
+def to_np_dict(item: t.Dict[str, t.Any]) -> t.Dict[str, np.ndarray]:
+    """Converts values in a dict to numpy arrays."""
+    return {k: np.asarray(v) for k, v in item.items()}
+
+
+def create_sq_ds_from_hf(
+    hf_ds: str,
+    hf_split: str,
+    to_np_dict: t.Callable[[t.Any], t.Dict[str, np.ndarray]],
+    shard_size: int,
+    squirrel_url: str = None,
+    **hf_kwargs,
+) -> None:
+    """Creates a squirrel Messagepack dataset from a Huggingface dataset.
+
+    Args:
+        hf_ds (str): The Huggingface dataset name.
+        hf_split (str): The Huggingface split name.
+        to_np_dict (t.Callable[[t.Any], t.Dict[str, np.ndarray]]): Function converting data to numpy arrays.
+        shard_size (int): Shard size of the squirrel dataset.
+        squirrel_url (str, optional): Where to save squirrel dataset. Defaults to None.
+    """
+
+    it = HuggingfaceDriver(hf_ds, **hf_kwargs).get_iter(hf_split)
+
+    # create a default path from hf data
+    if squirrel_url is None:
+        url = Path(hf_ds) / Path(hf_split)
+        url.mkdir(parents=True, exist_ok=True)
+        url = str(url)  # squirrel wants str
+
+    store = MessagepackDriver(url).store
+
+    print("creating squirrel dataset under", url, "...")
+    (
+        it.tqdm()  # prints "12345it [00:08, 5810.89it/s]"
+        .map(to_np_dict)
+        .batched(shard_size, drop_last_if_not_full=False)
+        .map(store.set)
+        .join()  # ensures direct iteration over stream
+    )
+
+
+if __name__ == "__main__":
+    create_sq_ds_from_hf("cifar100", to_np_dict=to_np_dict, shard_size=10000, hf_split="train")

docs/index.rst

@@ -37,6 +37,7 @@ Squirrel Documentation
    integration/spark.rst
    integration/dask.rst
    integration/squirrel_datasets_drivers.rst
+   examples/squirrel_dali.ipynb
 
 .. toctree::
    :maxdepth: 1

requirements.doc.in

@@ -6,5 +6,6 @@ sphinx_versions
 sphinx_rtd_theme
 sphinx-autoapi
 sphinxcontrib-mermaid
+myst_nb    # to use jupyter notebooks in docs
 Jinja2<3.1 # remove when sphinx is updated to 4.4, https://github.com/sphinx-doc/sphinx/issues/1029
 click<8.1  # remove when update docs, https://github.com/streamlit/streamlit/issues/4555