TensorFlow Dataset随机取样本内存溢出问题求助

方案1：直接通过Dataset操作实现随机选取 #

无需将整个数据集转为列表，利用unbatch() + shuffle() + take()组合直接获取随机样本，全程基于TensorFlow的Dataset管道操作，内存占用极低：

# 解除分批，打乱数据集（buffer_size设为总样本数保证完全随机），取10个随机样本
random_samples = test_data.unbatch().shuffle(buffer_size=25250).take(10)

# 遍历样本进行评估
for img, label in random_samples:
    # 执行模型预测或可视化等操作
    pred = model.predict(tf.expand_dims(img, axis=0), verbose=0)
    # 后续逻辑处理...

注意：buffer_size设置为测试集总样本数（25250）是为了让shuffle操作能从全量样本中随机选取，避免因buffer过小导致随机度不足。如果内存紧张，也可以适当减小buffer_size，但随机效果会打折扣。

方案2：通过索引过滤选取指定随机样本 #

先给每个样本添加索引，再随机生成目标索引，最后过滤出对应样本，适合需要精准选取特定数量随机样本的场景：

import tensorflow as tf
import numpy as np

# 为每个样本添加索引
indexed_dataset = test_data.unbatch().enumerate()

# 随机生成10个不重复的目标索引
total_samples = 25250
selected_indices = np.random.choice(total_samples, size=10, replace=False)

# 过滤出选中索引对应的样本
selected_samples = indexed_dataset.filter(
    lambda idx, data: tf.reduce_any(tf.equal(idx, selected_indices))
)

# 遍历处理样本
for idx, (img, label) in selected_samples:
    # 评估逻辑...
    pass

方案3：低内存批量转列表（仅当必须存全量样本时使用） #

如果确实需要将所有样本转为列表，不要一次性unbatch()后转，而是分批遍历处理，每处理一批就手动释放内存：

import gc

images = []
labels = []

# 逐批遍历数据集
for batch_imgs, batch_labels in test_data:
    # 拆分批次中的单个样本并添加到列表
    for img, label in zip(batch_imgs.numpy(), batch_labels.numpy()):
        images.append(img)
        labels.append(label)
    # 手动触发垃圾回收，释放批次内存
    gc.collect()

说明：直接转图片列表崩溃的原因是25250张224×224×3的float32图片约占15GB内存，远超Colab的常规内存上限，分批处理能避免一次性加载全量数据。

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