Hi,
Recently I tried to train an image captioning flax model with TPU, and I found it is very slow → each batch (256 images) requires 26 seconds. I finally found the cause:
My dataset have a column named “image” and it stores image array (being cached).
However, operations like batch = dataset[idx], where idx is a list of length 256 will take 15 seconds, if the image size is (256 , 256). If I store the flattened image array (i.e. 1D array as length 256 * 256 * 3), it will be faster.
The provided script demonstrate this situation, with image/batch size 128.
I am wondering if I am doing something wrong, and if there is a better way to handle image datasets.
import numpy as np
import time

from datasets import load_dataset, Dataset


ds = load_dataset("cifar10")
train_ds = ds["train"]
# Take 512 examples
train_ds = train_ds.select(range(512))


# `_images` has shape (batch_size, 32, 32, 3)
# Let's make large images (batch_size, 32 * FACTOR, 32 * FACTOR, 3)
FACTOR = 4


def preprocess_function(examples):

    inputs = {}
    images = examples["img"]

    # convert to np.float3f2 array
    _images = np.array(images, dtype=np.float32)

    # Make image larger
    batch_size, image_size = _images.shape[0:2]

    image_size = image_size * FACTOR

    # use (H, W, C) format
    _images = np.concatenate([_images] * FACTOR ** 2, axis=0).reshape((batch_size, image_size, image_size, 3))
    # flatten images to 1D -> `data_loader` will be faster
    # _images = _images.reshape((batch_size, image_size * image_size * 3))

    inputs["image"] = [x for x in _images]
    inputs["label"] = examples["label"]

    return inputs


times_1 = []
times_2 = []


def data_loader(rng: np.random.Generator, dataset: Dataset, batch_size: int, shuffle: bool = False):
    """
    Returns batches of size `batch_size` from truncated `dataset`, sharded over all local devices.
    Shuffle batches if `shuffle` is `True`.
    """

    steps_per_epoch = len(dataset) // batch_size

    if shuffle:
        batch_idx = rng.permutation(len(dataset))
    else:
        batch_idx = range(len(dataset))

    batch_idx = batch_idx[: steps_per_epoch * batch_size]  # Skip incomplete batch.
    batch_idx = batch_idx.reshape((steps_per_epoch, batch_size))

    for idx in batch_idx:

        s = time.time()
        # (bs=128, img_size=128)
        #   - (H, W, C) -> takes 6.959 seconds
        #   - flatten   -> takes 2.407 seconds
        batch = dataset[idx]
        e = time.time()
        times_1.append(e-s)

        s = time.time()
        # (bs=128, img_size=128)
        #   - (H, W, C) -> takes 1.661 seconds
        #   - flatten   -> takes 0.673 seconds
        batch = {k: np.array(v) for k, v in batch.items()}
        e = time.time()
        times_2.append(e-s)

        yield batch


# The results are cached
train_ds = train_ds.map(
    preprocess_function,
    batched=True,
    batch_size=16,
    num_proc=2,
)

# Create sampling rng
input_rng = np.random.default_rng(seed=42)
train_loader = data_loader(input_rng, train_ds, batch_size=128, shuffle=True)

start = time.time()
for idx, batch in enumerate(train_loader):

    end = time.time()
    print(f"batch: {idx} | time: {end - start}")
    start = time.time()


print(f"Average times_1 in data_loader: {np.mean(times_1)}")
print(f"Average times_2 in data_loader: {np.mean(times_2)}")