Feeding data to fit()

Keras · Practical · 12 min

Real datasets do not fit in memory; Keras 3 accepts the data pipelines of every major framework.

fit(), evaluate() and predict() accept several kinds of input, regardless of backend:

  • NumPy arrays (or backend tensors): simplest; pass batch_size and let Keras shuffle each epoch.
  • A `tf.data.Dataset`: yields batches already; do not pass batch_size. Works with JAX and PyTorch backends too, since Keras converts each batch.
  • A PyTorch `DataLoader`: yields batches; again no batch_size.
  • A `keras.utils.PyDataset`: a Python class with __len__ (number of batches) and __getitem__(idx) returning one batch. Pass workers= and use_multiprocessing= to its constructor to load in parallel.

For generators or infinite datasets, set steps_per_epoch so Keras knows when an epoch ends.

Batches can be (x, y) or (x, y, sample_weight), and x/y can be dicts matching named inputs and outputs. class_weight={0: 1.0, 1: 5.0} up-weights rare classes.

Pick by need: NumPy for small data, tf.data for high-throughput pipelines (especially with the TF or JAX backend), DataLoader if you already have PyTorch datasets, and PyDataset for custom Python loading code without extra dependencies.

Example

import math
import numpy as np
import keras

class RandomBatches(keras.utils.PyDataset):
    def __init__(self, n, batch_size, **kwargs):
        super().__init__(**kwargs)           # e.g. workers=4, use_multiprocessing=True
        self.n, self.batch_size = n, batch_size

    def __len__(self):                       # batches per epoch
        return math.ceil(self.n / self.batch_size)

    def __getitem__(self, idx):              # one batch
        size = min(self.batch_size, self.n - idx * self.batch_size)
        x = np.random.rand(size, 4).astype("float32")
        return x, x.sum(axis=1, keepdims=True)

model = keras.Sequential([keras.Input((4,)), keras.layers.Dense(1)])
model.compile(optimizer="adam", loss="mse")
model.fit(RandomBatches(1000, 64), epochs=2, verbose=0)
# Also valid for fit(): NumPy arrays + batch_size, a tf.data.Dataset,
# or a torch DataLoader, with any backend.

Remember

  • NumPy + batch_size, tf.data, torch DataLoader or PyDataset
  • Batched sources must not get a batch_size argument
  • class_weight and sample_weight reweight examples

Watch out for

  • Passing batch_size together with a dataset that is already batched.

Check yourself

  1. What must PyDataset.__len__ return?

Sign in to answer, see the explanations, and get this lesson back as a quick review later.

Official docs: keras.io/guides/training_with_built_in_methods/