Grain (ML Data Library)

0.2.16 · active · verified Sat Apr 11

Grain is a Python library from Google designed for efficiently loading and transforming data, primarily for machine learning model training and evaluation, particularly with JAX. It emphasizes flexibility, speed, and determinism in data processing pipelines. The library is actively developed, currently at version 0.2.16, with frequent updates including new features, bug fixes, and deprecations.

Warnings

breaking Custom implementations of `RandomAccessDataSource` must now accept an `int` index in `__getitem__`. While legacy paths handling `SupportsIndex` still work at runtime, type checkers may flag errors. Switch to `int` for full compatibility.
Fix: Update `__getitem__` methods in custom `RandomAccessDataSource` implementations to accept an `int` argument: `def __getitem__(self, index: int):`
deprecated Support for Python 3.10 has been deprecated, and the library now requires Python >=3.11.
Fix: Upgrade your Python environment to 3.11 or newer.
deprecated Experimental APIs `grain.python.experimental.MultiprocessPrefetchIterDataset` and `grain.python.experimental.ConcatenateMapDataset` have been deprecated. Use their graduated versions `grain.IterDataset.mp_prefetch` and `grain.MapDataset.concatenate` respectively.
Fix: Migrate usage from `grain.python.experimental.MultiprocessPrefetchIterDataset` to `grain.IterDataset.mp_prefetch` and from `grain.python.experimental.ConcatenateMapDataset` to `grain.MapDataset.concatenate`.
gotcha When using Python multiprocessing for parallel data loading and transformations, all custom transformation functions (e.g., `MapTransform` subclasses) must be picklable. Non-picklable objects or closures can lead to errors during serialization.
Fix: Ensure that any custom transformations are defined at the top level of a module or as static methods/free functions, and avoid using complex closures or unpicklable objects within them.
gotcha Choose between `MapDataset` and `IterDataset` based on access patterns. `MapDataset` supports efficient random access and is suitable for debugging or when order-dependent operations are needed. `IterDataset` (often created via `MapDataset.to_iter_dataset()`) is designed for performant, sequential iteration, typically used for training loops, especially with prefetching.
Fix: For random access or debugging, use `grain.MapDataset`. For performance-critical iteration during training, convert to `grain.IterDataset` using `dataset.to_iter_dataset()`.

Install

pip install grain Install stable release

Imports

MapDataset

import grain
dataset = grain.MapDataset.source([...])

IterDataset
```
import grain
iter_dataset = dataset.to_iter_dataset()
```
MultiprocessPrefetchIterDataset and ConcatenateMapDataset were deprecated in 0.2.16; use `IterDataset.mp_prefetch` or `MapDataset.concatenate` instead.

Quickstart

This example demonstrates how to create a simple `MapDataset` from a list, apply common transformations like shuffling, mapping, and batching, and then iterate through the processed data. It showcases the declarative chaining API for data pipeline construction.

import grain

dataset = (
    grain.MapDataset.source([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
    .shuffle(seed=42) # Shuffles elements globally.
    .map(lambda x: x + 1) # Maps each element.
    .batch(batch_size=2) # Batches consecutive elements.
)

print("Processing dataset:")
for batch in dataset:
    print(batch)

view raw JSON →