Mini Batches Dataloading#
How to load mini batches?#
The examples that I’ve shown so far are just toy examples. Nobody in real life will train neural networks with the entire dataset in one go or iterate through each data instance in the dataset. You use mini batches.
Jax and the ecosystem of libaries were primarily built with doing computation (or using fancy words, doing machine learning research) in mind. As a result, Jax has nothing similar to Tensorflow Datasets or Pytorch data utils for structuring and loading data. Jax instead encourages to use PyTorch and TF for data loading.
As outlined in one of the official guides, you can just use a pytorch dataset, and then override the collate_fn in the pytorch dataloader to use it with jax arrays. This may sound easy but for complex data cases writing a new collate_fn can be daunting and prone to errors. You can also use TF datasets but, personally, I’m not a big fan of TF and setting it up has been akin to offering peace to some god of the pantheon these days. So sticking with Pytorch, there must be an easier way right?
Well turns out, someone has written a package named jax-dataloader, which can take care of the dataloading functionality for you. All you need is to define a torch dataset. (If you’re visiting the jax-dataloader repo, give a star, and also contribute, we need more work like this).
So in this notebook, I’ll define a torch dataset and show you how to use jax-dataloader with it. The dataset will come from Huggingface Datasets package.
from datasets import load_dataset
dataset_name = "mnist"
dataset = load_dataset(dataset_name, split="train")
dataset
Dataset({
features: ['image', 'label'],
num_rows: 60000
})
# let's inspect a bit
sample = dataset[0]
print(type(sample["image"]))
print(type(sample["label"]))
<class 'PIL.PngImagePlugin.PngImageFile'>
<class 'int'>
import numpy as np
print(np.array(sample["image"]).shape)
(28, 28)
Let’s define the torch dataset.
from torch.utils.data import Dataset
class MNISTDataset(Dataset):
def __init__(self, split) -> None:
self.dataset_name = "mnist"
self.split = split
self.data = load_dataset(self.dataset_name, split=self.split)
def __len__(self):
return len(self.data)
def __getitem__(self, index):
data_instance = self.data[index]
image = data_instance["image"]
label = data_instance["label"]
# the image is a pillow image so we'll have to convert it to
# a numpy array
image = np.array(image)
return image, label
# create a training dataset
train_set = MNISTDataset("train")
print(type(train_set[0][0]))
<class 'numpy.ndarray'>
Now for the dataloader.
import jax_dataloader as jdl
train_loader = jdl.DataLoader(dataset=train_set, backend="pytorch", batch_size=64, shuffle=True)
So there are a few important things to notice here. The API of jdl.Dataloader is quite similar to that of Pytorch Dataloader. However, there is a backend parameter. This is to tell the jdl.Dataloader that the dataset will come in a pytorch format.
# iterate through the loader
for batch in train_loader:
print(batch)
break # let's stop here!
[array([[[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
...,
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0]],
[[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
...,
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0]],
[[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
...,
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0]],
...,
[[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
...,
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0]],
[[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
...,
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0]],
[[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
...,
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0]]], dtype=uint8), array([7, 8, 8, 0, 8, 4, 3, 9, 4, 3, 0, 6, 2, 1, 9, 6, 6, 7, 6, 5, 3, 8,
1, 1, 0, 3, 5, 4, 8, 3, 1, 4, 2, 2, 1, 0, 0, 6, 7, 5, 9, 5, 7, 1,
5, 7, 0, 7, 0, 7, 5, 0, 0, 6, 3, 6, 4, 0, 8, 1, 6, 9, 5, 1])]
What if I don’t want to use torch dataset?#
Good question. For simpler dataset structures you may not need something as expressive as a Dataset class. You can instead use the ArrayDataset type from jax-dataloader.
images = np.array(
[np.array(d["image"]) for d in dataset]
)
labels = np.array([d["label"] for d in dataset])
arr_dataset = jdl.ArrayDataset(images, labels)
# dataloader
arr_loader = jdl.DataLoader(dataset=arr_dataset, backend="jax", batch_size=64, shuffle=True)
# iterate
for batch in arr_loader:
print(batch)
break
(array([[[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
...,
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0]],
[[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
...,
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0]],
[[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
...,
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0]],
...,
[[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
...,
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0]],
[[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
...,
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0]],
[[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
...,
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0]]], dtype=uint8), array([9, 0, 4, 8, 4, 4, 5, 8, 9, 2, 7, 2, 5, 5, 9, 1, 2, 2, 2, 4, 9, 7,
6, 8, 0, 9, 1, 4, 1, 6, 0, 0, 3, 0, 3, 7, 2, 7, 4, 8, 3, 7, 7, 1,
6, 8, 3, 5, 3, 2, 5, 4, 6, 1, 7, 3, 6, 5, 5, 4, 3, 5, 1, 9]))
So what has changed here? The backend param. It’s no longer a pytorch dataset so you’ll mention jax as a backend. There’s another thing you can do with the ArrayDataset. Say, your dataset was originally given as pytorch tensors. How can you use them here? Just change the backend param.
import torch
images_tensor = torch.from_numpy(images)
labels_tensor = torch.from_numpy(labels)
ds = jdl.ArrayDataset(images_tensor, labels_tensor)
loader = jdl.DataLoader(dataset=ds, backend="pytorch", batch_size=64, shuffle=True)
for batch in loader:
print(batch)
break
[array([[[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
...,
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0]],
[[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
...,
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0]],
[[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
...,
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0]],
...,
[[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
...,
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0]],
[[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
...,
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0]],
[[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
...,
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0]]], dtype=uint8), array([9, 4, 1, 0, 4, 3, 1, 1, 7, 0, 9, 8, 5, 9, 7, 2, 9, 1, 2, 9, 2, 4,
4, 1, 4, 1, 3, 4, 7, 3, 2, 0, 5, 8, 9, 8, 2, 3, 1, 0, 1, 8, 7, 5,
8, 2, 9, 0, 0, 9, 2, 8, 4, 6, 0, 9, 5, 4, 5, 3, 2, 5, 5, 4])]