TensorFlow
next_batch
dataset
machine learning
Python

TensorFlow how is dataset.train.next_batch defined?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

dataset.train.next_batch(...) is not a core TensorFlow 2 API. It came from older TensorFlow 1 tutorials, especially the MNIST helper code in input_data.py, where a small Python DataSet class kept arrays in memory and returned slices batch by batch.

What next_batch Did in Old Tutorials

The historical method tracked three main pieces of state:

  • the current index within the dataset
  • the completed epoch count
  • whether to shuffle when a new epoch started

A simplified version looks like this:

python
1import numpy as np
2
3class DataSet:
4    def __init__(self, images, labels):
5        self.images = images
6        self.labels = labels
7        self.num_examples = len(images)
8        self.epochs_completed = 0
9        self.index_in_epoch = 0
10
11    def next_batch(self, batch_size, shuffle=True):
12        start = self.index_in_epoch
13
14        if start == 0 and self.epochs_completed == 0 and shuffle:
15            perm = np.random.permutation(self.num_examples)
16            self.images = self.images[perm]
17            self.labels = self.labels[perm]
18
19        if start + batch_size > self.num_examples:
20            self.epochs_completed += 1
21            rest = self.num_examples - start
22            images_rest = self.images[start:self.num_examples]
23            labels_rest = self.labels[start:self.num_examples]
24
25            if shuffle:
26                perm = np.random.permutation(self.num_examples)
27                self.images = self.images[perm]
28                self.labels = self.labels[perm]
29
30            self.index_in_epoch = batch_size - rest
31            end = self.index_in_epoch
32            images_new = self.images[0:end]
33            labels_new = self.labels[0:end]
34            return np.concatenate((images_rest, images_new)), np.concatenate((labels_rest, labels_new))
35
36        self.index_in_epoch += batch_size
37        end = self.index_in_epoch
38        return self.images[start:end], self.labels[start:end]

That is the essential idea: slice arrays, shuffle at epoch boundaries, and wrap around when the batch crosses the end of the data.

Why You Rarely See It Now

Modern TensorFlow uses tf.data.Dataset, which separates data input pipelines from ad hoc Python helper classes.

python
1import tensorflow as tf
2
3images = tf.random.uniform((100, 28, 28))
4labels = tf.random.uniform((100,), maxval=10, dtype=tf.int32)
5
6dataset = tf.data.Dataset.from_tensor_slices((images, labels))
7dataset = dataset.shuffle(100).batch(32).prefetch(tf.data.AUTOTUNE)
8
9for batch_images, batch_labels in dataset.take(1):
10    print(batch_images.shape, batch_labels.shape)

The TensorFlow tf.data guide documents Dataset.batch() and related transformations such as shuffle() and prefetch() as the modern batching path.

Conceptual Mapping from Old to New

The rough translation is:

  • 'next_batch(size) becomes Dataset.batch(size)'
  • shuffle-on-epoch becomes Dataset.shuffle(...).repeat()
  • manual Python slicing becomes a composable input pipeline

That makes training code more scalable and more compatible with GPUs, TPUs, and distributed execution.

If you want an endless training stream similar to repeated next_batch calls across epochs, combine batching with repeat() in the tf.data pipeline. That moves epoch rollover logic out of your Python helper class and into the data pipeline itself.

Why the Old Helper Existed

Early TensorFlow tutorials optimized for simple, readable notebook code. Keeping MNIST in NumPy arrays and returning one batch at a time was easy to explain.

That design was fine for small examples, but it was not the long-term input pipeline model TensorFlow standardized on.

It also meant the batching logic was visible Python code, which made tutorials approachable but limited for larger datasets and production training jobs.

Common Pitfalls

The most common mistake is searching for dataset.train.next_batch in modern TensorFlow APIs and assuming it still exists. It does not as a standard TF 2 workflow.

Another issue is copying old TF 1 tutorial code into a TF 2 project and then trying to mix it with eager execution and tf.data.

A third pitfall is forgetting that the old helper was just Python array slicing with bookkeeping, not a magical TensorFlow primitive.

Summary

  • 'dataset.train.next_batch came from older TensorFlow tutorial helper code, not the modern core API.'
  • It returned slices of in-memory arrays and handled epoch rollover and optional shuffling.
  • The modern replacement is tf.data.Dataset.shuffle(...).batch(...).prefetch(...).
  • Understanding the old method is useful when reading TF 1 tutorials.
  • For new code, use tf.data instead of re-creating next_batch by hand.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.