TensorFlow
minibatch
numpy
large arrays
data processing

Tensorflow create minibatch from numpy array 2 GB

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

When a NumPy array is around 2 GB, the real problem is not minibatching itself. The real problem is avoiding unnecessary copies and avoiding a training loop that tries to move the whole array through Python on every step. TensorFlow's tf.data pipeline is the right tool because it batches lazily and keeps the input path closer to the runtime.

Avoid Feeding the Whole Array Repeatedly

A common beginner pattern is calling Session.run(..., feed_dict=...) or repeatedly slicing Python arrays inside a training loop. With very large arrays, that creates extra copies, heavy Python overhead, and unpredictable memory pressure.

If the array already fits in RAM and is a normal numpy.ndarray, tf.data.Dataset.from_tensor_slices is the cleanest starting point.

python
1import numpy as np
2import tensorflow as tf
3
4x = np.random.rand(10000, 128).astype("float32")
5y = np.random.randint(0, 2, size=(10000,)).astype("int32")
6
7dataset = tf.data.Dataset.from_tensor_slices((x, y))
8dataset = dataset.shuffle(10000).batch(64).prefetch(tf.data.AUTOTUNE)
9
10for batch_x, batch_y in dataset.take(1):
11    print(batch_x.shape, batch_y.shape)

This does not train on the entire array at once. It exposes minibatches as the dataset is consumed.

Use Memory Mapping When the Array Is Too Big for Comfortable RAM Use

If 2 GB is technically loadable but leaves the system under memory pressure, store the data as a memory-mapped NumPy array and batch from that. np.memmap lets the operating system page data in as needed instead of forcing one giant in-memory allocation.

python
1import numpy as np
2import tensorflow as tf
3
4x_mem = np.memmap("features.dat", dtype="float32", mode="r", shape=(500000, 128))
5y_mem = np.memmap("labels.dat", dtype="int32", mode="r", shape=(500000,))
6
7dataset = tf.data.Dataset.from_tensor_slices((x_mem, y_mem))
8dataset = dataset.batch(128).prefetch(tf.data.AUTOTUNE)

This is a strong pattern when the bottleneck is memory footprint rather than raw model compute.

Generator Pipelines Are Useful for More Complex Storage Layouts

If the data is not stored as one clean contiguous array, use from_generator so each batch or sample is yielded on demand.

python
1import numpy as np
2import tensorflow as tf
3
4
5def batch_generator(x, y, batch_size):
6    for start in range(0, len(x), batch_size):
7        end = start + batch_size
8        yield x[start:end], y[start:end]
9
10
11x = np.random.rand(10000, 128).astype("float32")
12y = np.random.randint(0, 2, size=(10000,)).astype("int32")
13
14dataset = tf.data.Dataset.from_generator(
15    lambda: batch_generator(x, y, 64),
16    output_signature=(
17        tf.TensorSpec(shape=(None, 128), dtype=tf.float32),
18        tf.TensorSpec(shape=(None,), dtype=tf.int32),
19    ),
20)
21
22for batch_x, batch_y in dataset.take(1):
23    print(batch_x.shape, batch_y.shape)

This is slightly more Python-heavy than from_tensor_slices, but it gives you control when the source data is irregular or lives in custom binary storage.

Shuffle Deliberately

With very large arrays, a full in-memory shuffle may be expensive. TensorFlow's shuffle buffer gives a practical compromise between randomness and memory usage.

python
dataset = tf.data.Dataset.from_tensor_slices((x, y))
dataset = dataset.shuffle(buffer_size=10000)
dataset = dataset.batch(64)

The larger the buffer, the closer the shuffle is to a full permutation. The smaller the buffer, the lower the memory overhead.

Integrate Directly with Keras Training

Once the dataset yields correctly shaped minibatches, pass it straight into model.fit.

python
1from tensorflow import keras
2
3model = keras.Sequential([
4    keras.layers.Input(shape=(128,)),
5    keras.layers.Dense(64, activation="relu"),
6    keras.layers.Dense(1, activation="sigmoid"),
7])
8
9model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])
10model.fit(dataset, epochs=3)

This is cleaner than manually slicing batches in Python loops, and it gives TensorFlow room to pipeline data loading with compute.

When a Single NumPy Array Is the Wrong Storage Format

If the dataset is much larger than memory or must be shared across machines, a single .npy or in-memory array may be the wrong long-term format. In those cases, consider sharded files such as TFRecords or on-disk chunked formats that stream naturally into tf.data.

Minibatching from a 2 GB NumPy array is possible. It is just not always the best architecture if the dataset keeps growing.

Common Pitfalls

A common mistake is creating minibatches by copying slices into new Python lists every training step. That wastes memory bandwidth and makes the CPU the bottleneck.

Another mistake is assuming that because the array fits in RAM once, every downstream step is cheap. Extra copies during shuffle, cast, or feed operations can still exhaust memory.

People also often overlook prefetch, which means data loading and model execution happen strictly in sequence.

Finally, if performance is poor, measure whether the problem is the model, the storage format, or the Python input path before blaming TensorFlow itself.

Summary

  • Use tf.data to minibatch large NumPy arrays lazily.
  • 'from_tensor_slices is the cleanest path when the array is already in memory.'
  • Use np.memmap when the array is too large for comfortable RAM use.
  • Prefer dataset pipelines over manual Python batch slicing in training loops.
  • If the dataset keeps growing, consider streaming-friendly storage instead of one giant array.

Related reading
Course
Intermediate
27 lessons
15 hours
DSA Fundamentals

Master algorithmic patterns and data structures through hands-on LeetCode-style problems - from arrays and hashing to dynamic programming and advanced graphs.

View the 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.