Simple usage example? (eg. MNIST)

[ameya98]

Hi, thanks for making this! Is there a simple training example, perhaps with MNIST or another small dataset?

[pablo2909]

+1 , that would be really cool !

[GrantMcConachie]

Any updates on this?

[simon-bachhuber]

Maybe it helps:

from dataclasses import dataclass

import grain.python as pygrain


@dataclass
class MyDataset:
    data: list

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx: int):
        return self.data[idx]


data_source = MyDataset(list(range(6)))

sampler = pygrain.IndexSampler(
    len(data_source),
    shuffle=True,
    seed=2,
    shard_options=pygrain.NoSharding(),
    num_epochs=3,
)

batch_size = 3
dl = pygrain.DataLoader(
    data_source=data_source, sampler=sampler, operations=[pygrain.Batch(batch_size)]
)

for sample in dl:
    print(sample)

"""Prints:
[0 4 1]
[2 3 5]
[4 3 0]
[1 2 5]
[5 0 4]
[3 2 1]
"""

[cnguyen10]

There is another example of grain using with tensorflow_datasets (or tfds) here: TFDS with Jax.

[selamw1]

This new tutorial explores different data loading strategies including Grain for a simple image classification task based on the MNIST dataset.