TensorFlow
MNIST
next_batch
machine learning
tutorials

Where does next_batch in the TensorFlow tutorial batch_xs, batch_ys mnist.train.next_batch100 come from?

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

In the original TensorFlow 1.x MNIST tutorial, mnist.train.next_batch(100) returns a batch of 100 training images and their labels. This method comes from tensorflow.examples.tutorials.mnist.input_data, which provides a DataSet class with a built-in next_batch() method that handles shuffling, batching, and epoch cycling. In TensorFlow 2.x, this tutorial API is deprecated — the modern equivalent is tf.data.Dataset with .batch() and .shuffle().

The Original TF 1.x Code

python
1# TensorFlow 1.x tutorial code
2from tensorflow.examples.tutorials.mnist import input_data
3
4mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)
5
6# mnist.train — 55,000 training examples
7# mnist.test — 10,000 test examples
8# mnist.validation — 5,000 validation examples
9
10for i in range(1000):
11    batch_xs, batch_ys = mnist.train.next_batch(100)
12    # batch_xs: shape (100, 784) — 100 images, 28x28 pixels flattened
13    # batch_ys: shape (100, 10)  — 100 one-hot labels
14    sess.run(train_step, feed_dict={x: batch_xs, y_: batch_ys})

next_batch(100) returns a random sample of 100 images from the training set. Each call returns a different batch. When all examples have been used, it reshuffles and starts a new epoch.

How next_batch Works Internally

The DataSet class maintained an internal index and shuffled indices:

python
1# Simplified version of what next_batch did internally
2class DataSet:
3    def __init__(self, images, labels):
4        self._images = images
5        self._labels = labels
6        self._num_examples = images.shape[0]
7        self._index_in_epoch = 0
8        self._epochs_completed = 0
9
10    def next_batch(self, batch_size):
11        start = self._index_in_epoch
12        self._index_in_epoch += batch_size
13
14        if self._index_in_epoch > self._num_examples:
15            # Epoch completed — shuffle and reset
16            self._epochs_completed += 1
17            perm = numpy.arange(self._num_examples)
18            numpy.random.shuffle(perm)
19            self._images = self._images[perm]
20            self._labels = self._labels[perm]
21            start = 0
22            self._index_in_epoch = batch_size
23
24        end = self._index_in_epoch
25        return self._images[start:end], self._labels[start:end]

TF 2.x: Modern Equivalent with tf.data

python
1import tensorflow as tf
2
3# Load MNIST dataset
4(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
5
6# Preprocess
7x_train = x_train.reshape(-1, 784).astype("float32") / 255.0
8y_train = tf.keras.utils.to_categorical(y_train, 10)
9
10# Create batched dataset
11train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
12train_dataset = train_dataset.shuffle(buffer_size=10000).batch(100)
13
14# Training loop
15model = tf.keras.Sequential([
16    tf.keras.layers.Dense(128, activation="relu", input_shape=(784,)),
17    tf.keras.layers.Dense(10, activation="softmax")
18])
19model.compile(optimizer="adam", loss="categorical_crossentropy", metrics=["accuracy"])
20
21# Option 1: fit with dataset
22model.fit(train_dataset, epochs=5)
23
24# Option 2: manual loop
25for epoch in range(5):
26    for batch_xs, batch_ys in train_dataset:
27        model.train_on_batch(batch_xs, batch_ys)

Key Differences: next_batch vs tf.data.Dataset

python
1# next_batch (TF 1.x)
2batch_xs, batch_ys = mnist.train.next_batch(100)
3# - Eager: returns numpy arrays immediately
4# - Shuffles at epoch boundary
5# - Simple but not optimized for large datasets
6
7# tf.data.Dataset (TF 2.x)
8dataset = tf.data.Dataset.from_tensor_slices((x, y))
9dataset = dataset.shuffle(10000).batch(100).prefetch(tf.data.AUTOTUNE)
10# - Lazy: builds a data pipeline
11# - Shuffles continuously with buffer
12# - Supports prefetching, parallel loading, caching
Featurenext_batch (TF 1.x)tf.data.Dataset (TF 2.x)
ShufflingFull reshuffle at epoch endRolling buffer shuffle
BatchingReturns NumPy arraysReturns tf.Tensor
PerformanceNo prefetchingSupports prefetch/cache
Large datasetsMust fit in memorySupports generators, files
API statusDeprecatedCurrent standard

Loading MNIST in TF 2.x

python
1import tensorflow as tf
2
3# Built-in Keras dataset
4(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
5print(x_train.shape)  # (60000, 28, 28)
6print(y_train.shape)  # (60000,) — integer labels, not one-hot
7
8# For one-hot labels
9y_train_onehot = tf.keras.utils.to_categorical(y_train, 10)
10
11# Full pipeline
12train_ds = (
13    tf.data.Dataset.from_tensor_slices((x_train, y_train))
14    .map(lambda x, y: (tf.cast(x, tf.float32) / 255.0, y))
15    .shuffle(10000)
16    .batch(128)
17    .prefetch(tf.data.AUTOTUNE)
18)
19
20model.fit(train_ds, epochs=10)

Common Pitfalls

  • Using the deprecated input_data module in TF 2.x: tensorflow.examples.tutorials.mnist was removed in TensorFlow 2.0. Use tf.keras.datasets.mnist.load_data() instead. Installing tensorflow-datasets is another option for more datasets.
  • Confusing next_batch shuffling with tf.data shuffling: next_batch reshuffled the entire dataset at each epoch boundary. tf.data.Dataset.shuffle(buffer_size) maintains a rolling buffer — only buffer_size elements are shuffled at a time. Set buffer_size equal to the dataset size for a full shuffle.
  • Forgetting to normalize pixel values: MNIST pixel values range from 0 to 255. Both the old tutorial and keras.datasets.mnist return raw integers. Always divide by 255.0 to normalize to [0, 1] before training.
  • Using feed_dict in TF 2.x: TF 2.x uses eager execution by default. sess.run() and feed_dict are TF 1.x patterns. Use model.fit(), model.train_on_batch(), or tf.GradientTape for TF 2.x training loops.
  • Not using prefetch in the data pipeline: Without prefetch(tf.data.AUTOTUNE), the GPU sits idle while the CPU prepares the next batch. Always add .prefetch() at the end of a tf.data pipeline for overlapped data loading and training.

Summary

  • mnist.train.next_batch(100) comes from TensorFlow 1.x's input_data module — it returns random batches with automatic shuffling and epoch tracking
  • This API was deprecated and removed in TensorFlow 2.0
  • The modern replacement is tf.data.Dataset.from_tensor_slices().shuffle().batch()
  • Use tf.keras.datasets.mnist.load_data() to load MNIST in TF 2.x
  • Add .prefetch(tf.data.AUTOTUNE) to the pipeline for optimal GPU utilization
  • Use model.fit(dataset) instead of sess.run() with feed_dict for TF 2.x training

Course illustration
Course illustration

All Rights Reserved.