Introduction

Feeding .npy files into a TensorFlow pipeline is completely workable, but the best approach depends on how many files you have and whether each file fits comfortably in memory. For small or medium arrays, loading with NumPy and then building a tf.data.Dataset is simple; for larger collections of .npy files, you usually want lazy loading and parallel mapping.

If the Array Fits in Memory

The simplest path is:

1. load the .npy file with NumPy
2. create a dataset from the loaded array
3. batch and prefetch

1import numpy as np
2import tensorflow as tf
3
4features = np.load("features.npy")
5labels = np.load("labels.npy")
6
7dataset = tf.data.Dataset.from_tensor_slices((features, labels))
8dataset = dataset.shuffle(1000).batch(32).prefetch(tf.data.AUTOTUNE)
9
10for batch_x, batch_y in dataset.take(1):
11    print(batch_x.shape, batch_y.shape)

This is the cleanest solution when both arrays fit in RAM and training simplicity matters more than elaborate I/O optimization.

Use Memory Mapping for Larger Arrays

If the array is large, np.load(..., mmap_mode="r") can reduce memory pressure by reading lazily from disk-backed memory maps.

1import numpy as np
2import tensorflow as tf
3
4features = np.load("features.npy", mmap_mode="r")
5labels = np.load("labels.npy", mmap_mode="r")
6
7dataset = tf.data.Dataset.from_tensor_slices((features, labels))
8dataset = dataset.batch(64).prefetch(tf.data.AUTOTUNE)

This still looks simple in code, but it avoids fully materializing the array in process memory up front.

Memory mapping works best when:

• the file format is already stable
• you need random indexed access
• the training host can stream from disk efficiently

If You Have Many .npy Files

When data is split across many .npy files, a common pattern is to list the files and load them lazily through tf.numpy_function.

1import numpy as np
2import tensorflow as tf
3
4def load_npy(path):
5    path = path.decode("utf-8")
6    array = np.load(path).astype("float32")
7    return array
8
9file_ds = tf.data.Dataset.list_files("data/*.npy", shuffle=True)
10
11dataset = file_ds.map(
12    lambda path: tf.numpy_function(load_npy, [path], tf.float32),
13    num_parallel_calls=tf.data.AUTOTUNE,
14)
15
16dataset = dataset.prefetch(tf.data.AUTOTUNE)

This is flexible, but it comes with an important catch: TensorFlow loses static shape information when numpy_function is involved.

Restore Shapes After numpy_function

If you know the array shape, set it explicitly after loading so the rest of the model pipeline can reason about tensor shapes correctly.

1import numpy as np
2import tensorflow as tf
3
4def load_image(path):
5    path = path.decode("utf-8")
6    array = np.load(path).astype("float32")
7    return array
8
9def tf_load_image(path):
10    tensor = tf.numpy_function(load_image, [path], tf.float32)
11    tensor.set_shape((128, 128, 3))
12    return tensor
13
14dataset = tf.data.Dataset.list_files("images/*.npy")
15dataset = dataset.map(tf_load_image, num_parallel_calls=tf.data.AUTOTUNE)
16dataset = dataset.batch(16).prefetch(tf.data.AUTOTUNE)

Without set_shape, downstream layers may receive tensors with unknown dimensions and fail later in confusing ways.

Pair Features and Labels Cleanly

If features and labels live in separate files, keep the mapping explicit. One approach is to store paired paths:

1import tensorflow as tf
2
3feature_paths = tf.constant(["x1.npy", "x2.npy"])
4label_paths = tf.constant(["y1.npy", "y2.npy"])
5
6path_ds = tf.data.Dataset.from_tensor_slices((feature_paths, label_paths))

Then load each side with a mapping function. The important part is that the file pairing remains deterministic and auditable.

When .npy Stops Being the Best Format

.npy is convenient, but not always ideal for large-scale training. If the pipeline becomes more production-oriented, formats such as TFRecord or Parquet may be a better fit because they integrate more naturally with streaming and distributed input pipelines.

That does not mean .npy is wrong. It just means:

• '.npy is great for experiments and simpler array persistence'
• other formats may scale better for long-term training systems

The right choice depends on the size and lifespan of the pipeline.

Common Pitfalls

• Loading every .npy file eagerly into RAM when the dataset is too large for comfortable memory use.
• Using tf.numpy_function without restoring static shapes afterward.
• Assuming many small .npy files will always perform well without measuring file-system overhead.
• Forgetting to cast NumPy arrays to the dtype the model actually expects.
• Forcing .npy into a large distributed training setup when a more streaming-friendly format would be easier to scale.

Summary

• If the arrays fit in memory, load them with NumPy and wrap them with tf.data.Dataset.from_tensor_slices.
• Use mmap_mode="r" when arrays are large and you want lighter memory usage.
• For many .npy files, load lazily through Dataset.list_files and tf.numpy_function.
• Restore tensor shapes after numpy_function so the rest of the pipeline stays well-defined.
• '.npy is excellent for many experimental workflows, but it is not always the best long-term training format.'