如何直接从115GB .npz文件创建tf.data.Dataset以避免大内存占用？

可以绕过一次性加载整个NumPy数组的步骤，核心思路是利用tf.data的流式读取能力，结合npz文件的结构特性分批次加载数据，避免占用全部内存。

方法1：用tf.data.Dataset.from_generator流式读取 #

npz本质是压缩归档文件，可通过mmap_mode='r'模式打开文件（仅将文件映射到内存，而非一次性加载全部内容），再用Python生成器分批次读取数据，最后包装成tf.data数据集。

示例代码：

import numpy as np
import tensorflow as tf

def npz_batch_generator(npz_path, batch_size=32):
    # 以内存映射模式打开npz，不加载全部数据
    with np.load(npz_path, mmap_mode='r') as npz_data:
        # 假设npz中包含'features'和'labels'两个数组，根据你的实际键名调整
        features = npz_data['features']
        labels = npz_data['labels']
        total_samples = features.shape[0]
        
        # 分批次生成数据
        for start_idx in range(0, total_samples, batch_size):
            end_idx = min(start_idx + batch_size, total_samples)
            # 仅加载当前批次的切片数据
            batch_feat = features[start_idx:end_idx]
            batch_lbl = labels[start_idx:end_idx]
            yield (batch_feat, batch_lbl)

# 创建数据集，需指定输出签名匹配数据形状和类型
dataset = tf.data.Dataset.from_generator(
    generator=npz_batch_generator,
    args=['your_115gb_file.npz'],
    output_signature=(
        tf.TensorSpec(shape=(None, *features.shape[1:]), dtype=tf.float32),
        tf.TensorSpec(shape=(None, *labels.shape[1:]), dtype=tf.int32)
    )
)

# 后续可添加预处理、打乱、预取等优化操作
dataset = dataset.shuffle(buffer_size=1000).prefetch(tf.data.AUTOTUNE)

方法2：拆分npz为小文件后并行读取 #

如果流式读取的性能达不到要求，可以先把大npz拆分成多个小npz文件，再通过tf.data批量读取：

import tensorflow as tf
import numpy as np

def load_single_npz(file_path):
    # 转换TensorFlow字符串路径为Python字符串
    path_str = file_path.numpy().decode('utf-8')
    with np.load(path_str) as data:
        return data['features'], data['labels']

# 列出所有拆分后的小npz文件
file_dataset = tf.data.Dataset.list_files('split_npz/*.npz')
# 并行读取每个小文件
dataset = file_dataset.map(
    lambda x: tf.py_function(load_single_npz, [x], [tf.float32, tf.int32]),
    num_parallel_calls=tf.data.AUTOTUNE
)
# 打散并重新分批次
dataset = dataset.unbatch().batch(32).prefetch(tf.data.AUTOTUNE)

注意事项 #

• 确保npz内的数组是连续存储的，否则内存映射模式的读取效率会下降。
• 输出签名（output_signature）必须和你的数据形状、数据类型严格匹配，否则会影响TensorFlow的图优化。
• 如果npz内包含多个独立数组（非样本维度拆分），需要调整生成器逻辑，按需读取对应数组。

内容的提问来源于stack exchange，提问作者Aspir