加载20GB NPZ文件内存超限（Colab Pro）的解决方案求助

你的问题核心是np.load()默认会把整个NPZ文件的数组一次性加载到内存，哪怕后续转成tf.data.Dataset，数组已经全部占用内存了，导致超过Colab Pro的51GB限制。以下是几个可行的解决方法：

方法1：使用内存映射模式加载NPZ #

利用np.load()的mmap_mode参数，将数组以内存映射的方式加载，数据不会一次性全部读入内存，而是按需从磁盘读取：

with np.load("/content/drive/MyDrive/NPZ/train.npz", mmap_mode='r') as f:
    keys = list(f.keys())
    # 内存映射的数组可直接传入tf.data.Dataset
    train_dataset = tf.data.Dataset.from_tensor_slices((f[keys[0]], f[keys[1]]))

这种方式能大幅降低初始内存占用，因为数组数据实际仍存储在磁盘上，只有当需要访问时才会加载对应部分到内存。

方法2：拆分大NPZ为多个小文件 #

把20GB的NPZ拆分成多个1GB-2GB的小NPZ文件，再用tf.data.Dataset批量加载：

1. 拆分原NPZ（示例代码）：

with np.load("/content/drive/MyDrive/NPZ/train.npz") as f:
    data1 = f[keys[0]]
    data2 = f[keys[1]]
# 按每1000个样本为一批拆分
split_size = 1000
for i in range(0, len(data1), split_size):
    np.savez(f"/content/drive/MyDrive/NPZ/train_split_{i//split_size}.npz", 
             x=data1[i:i+split_size], y=data2[i:i+split_size])

2. 加载拆分后的文件：

import glob

file_pattern = "/content/drive/MyDrive/NPZ/train_split_*.npz"
files = glob.glob(file_pattern)

def load_npz_file(file_path):
    def _load(path):
        path_str = path.numpy().decode("utf-8")
        with np.load(path_str) as f:
            return f["x"], f["y"]
    return tf.py_function(_load, [file_path], (tf.float32, tf.int32))  # 替换为你的数据类型

train_dataset = tf.data.Dataset.from_tensor_slices(files)
train_dataset = train_dataset.interleave(
    lambda x: tf.data.Dataset.from_generator(
        lambda: load_npz_file(x),
        output_types=(tf.float32, tf.int32)
    ),
    num_parallel_calls=tf.data.AUTOTUNE
)

这样每次只会加载一个小NPZ文件到内存，避免一次性占用大量内存。

方法3：转换为TFRecord格式（推荐） #

TFRecord是TensorFlow原生的高效存储格式，支持流式读取，非常适合大数据集：

1. 将NPZ数据转换为TFRecord：

import tensorflow as tf

def _bytes_feature(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

def _int64_feature(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))

with np.load("/content/drive/MyDrive/NPZ/train.npz") as f:
    data1 = f[keys[0]]
    data2 = f[keys[1]]

output_path = "/content/drive/MyDrive/TFRecord/train.tfrecord"
with tf.io.TFRecordWriter(output_path) as writer:
    for x, y in zip(data1, data2):
        # 根据数据类型序列化，示例中x为视频帧数组，y为标签
        x_bytes = x.tobytes()
        feature = {
            "x": _bytes_feature(x_bytes),
            "y": _int64_feature(y)  # 替换为你的标签类型
        }
        example = tf.train.Example(features=tf.train.Features(feature=feature))
        writer.write(example.SerializeToString())

2. 加载TFRecord数据集：

def parse_example(example_proto):
    feature_description = {
        "x": tf.io.FixedLenFeature([], tf.string),
        "y": tf.io.FixedLenFeature([], tf.int64)
    }
    parsed_features = tf.io.parse_single_example(example_proto, feature_description)
    # 还原数组，示例中x为(frames, height, width, channels)的视频数据
    x = tf.io.decode_raw(parsed_features["x"], tf.float32)
    x = tf.reshape(x, (16, 224, 224, 3))  # 替换为你的数据形状
    y = parsed_features["y"]
    return x, y

train_dataset = tf.data.TFRecordDataset("/content/drive/MyDrive/TFRecord/train.tfrecord")
train_dataset = train_dataset.map(parse_example, num_parallel_calls=tf.data.AUTOTUNE)
train_dataset = train_dataset.batch(32)  # 设置批次大小

TFRecord会按需读取数据，不会一次性加载整个数据集到内存，同时还能提升数据读取效率。

方法4：临时应急——加载部分数据 #

如果不想修改存储格式，可临时只读取部分数据分批处理，但仅适合调试，不适合完整训练：

with np.load("/content/drive/MyDrive/NPZ/train.npz", mmap_mode='r') as f:
    keys = list(f.keys())
    # 只加载前10000个样本
    train_dataset = tf.data.Dataset.from_tensor_slices((f[keys[0]][:10000], f[keys[1]][:10000]))

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