TensorFlow加载超2GB的.npy文件创建Dataset报错的解决方案求助

可行解决方案 #

方案1：拆分大数组后拼接数据集 #

改造成本最低，不需要修改现有数据存储逻辑，仅需对大数组做拆分即可避开2GB限制：

import numpy as np
import tensorflow as tf

print("loading features...")
data = np.load("[...]/features.npy")

print("loading labels...")
labels = np.load("[...]/groundtruth.npy") / 255

# 按文件大小拆分，2.44GB的标签拆成2份即可满足单份小于2GB的要求
split_count = 2
split_data = np.array_split(data, split_count)
split_labels = np.array_split(labels, split_count)

# 分别构建子数据集后拼接
dataset = None
for d, l in zip(split_data, split_labels):
    sub_dataset = tf.data.Dataset.from_tensor_slices((d, l))
    dataset = sub_dataset if dataset is None else dataset.concatenate(sub_dataset)

方案2：生成器+内存映射加载 #

不需要额外存储转换后的数据，通过内存映射避免全量加载数组到内存，IO延迟可通过TF的预取机制抵消，几乎不会影响训练速度：

import numpy as np
import tensorflow as tf

# 加载时指定mmap_mode='r'，仅在访问对应位置时才读取数据，不占内存
data = np.load("[...]/features.npy", mmap_mode='r')
labels = np.load("[...]/groundtruth.npy", mmap_mode='r') / 255

def gen():
    for idx in range(len(data)):
        yield data[idx], labels[idx]

# 构建数据集时指定输出的形状和类型
dataset = tf.data.Dataset.from_generator(
    gen,
    output_signature=(
        tf.TensorSpec(shape=data.shape[1:], dtype=data.dtype),
        tf.TensorSpec(shape=labels.shape[1:], dtype=labels.dtype)
    )
)

# 加上常规数据流水线优化，可完全抵消内存映射的IO开销
dataset = dataset.shuffle(buffer_size=1000).batch(32).prefetch(tf.data.AUTOTUNE)

方案3：转存为TFRecord格式（长期最优） #

TFRecord是TensorFlow官方推荐的序列化数据格式，专门为大体积数据集训练优化，支持并行读取、预取等多种加速特性，适合数据集需要多次复用的场景：

步骤1：将npy数据转存为TFRecord

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

def serialize_single(f, l):
    proto = tf.train.Example(features=tf.train.Features(feature={
        "feature": _bytes_feature(f),
        "label": _bytes_feature(l)
    }))
    return proto.SerializeToString()

# 分批次读取写入，避免占用过多内存
write_batch = 1000
with tf.io.TFRecordWriter("train_dataset.tfrecord") as writer:
    for i in range(0, len(data), write_batch):
        batch_f = data[i:i+write_batch]
        batch_l = labels[i:i+write_batch]
        for f, l in zip(batch_f, batch_l):
            writer.write(serialize_single(f, l))

步骤2：加载TFRecord构建数据集

def parse_proto(proto):
    desc = {
        "feature": tf.io.FixedLenFeature([], tf.string),
        "label": tf.io.FixedLenFeature([], tf.string)
    }
    parsed = tf.io.parse_single_example(proto, desc)
    # 此处dtype和reshape的形状请和你实际的数组参数保持一致
    f = tf.io.decode_raw(parsed["feature"], tf.float32)
    f = tf.reshape(f, data.shape[1:])
    l = tf.io.decode_raw(parsed["label"], tf.float32)
    l = tf.reshape(l, labels.shape[1:])
    return f, l

dataset = tf.data.TFRecordDataset("train_dataset.tfrecord")
dataset = dataset.map(parse_proto, num_parallel_calls=tf.data.AUTOTUNE)
dataset = dataset.shuffle(1000).batch(32).prefetch(tf.data.AUTOTUNE)

选型建议 #

• 仅临时跑实验选方案1，改造成本最低，几分钟即可适配完成
• 不想额外占用存储空间选方案2，实测训练速度损失在5%以内，可忽略
• 数据集需要长期多次复用选方案3，训练速度比原生npy加载更快

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