使用tf.train.shuffle_batch处理大数据集时遇2GB张量报错求解

首先，咱们得先搞清楚你遇到这个错误的根本原因：

ValueError: Cannot create a tensor proto whose content is larger than 2GB.

当你设置enqueue_many=True时，tf.train.shuffle_batch会尝试把整个images["train"]和labels["train"]张量直接序列化到TensorFlow的计算图结构里。咱们算一下你的tiny-Imagenet训练集大小：88179张64×64×3的图片，按float32类型存储的话，总大小大概是88179 * 64 * 64 * 3 * 4 ≈ 4.3GB，远超TensorFlow对单个TensorProto的2GB限制，所以直接触发了错误。而CIFAR-10的张量总大小只有约0.5GB，没超限制，所以能正常运行。

接下来给你两个可行的解决方案，按推荐程度排序：

方案一：改用tf.data.Dataset API（推荐） #

TensorFlow 1.4已经支持tf.data API了，这个API专门为处理大数据集设计，从根源上避免了这种张量序列化的问题，而且代码更简洁灵活。

示例代码如下：

# 从内存中的张量创建数据集
train_dataset = tf.data.Dataset.from_tensor_slices((images["train"], labels["train"]))

# 打乱数据（buffer_size建议设为数据集大小的1/10到1/5，保证 shuffle 效果）
train_dataset = train_dataset.shuffle(buffer_size=50000, seed=self.seed)

# 分批，drop_remainder=False对应原来的allow_smaller_final_batch=True
train_dataset = train_dataset.batch(self.batch_size, drop_remainder=False)

# 创建可初始化迭代器
iterator = train_dataset.make_initializable_iterator()
x_train, y_train = iterator.get_next()

# 在会话中运行时，需要先初始化迭代器
with tf.Session() as sess:
    sess.run(iterator.initializer)
    # 后续的训练循环逻辑

这个方式不会把整个大张量序列化到计算图里，而是在运行时动态切片读取，完美避开2GB限制。

方案二：用slice_input_producer替代直接传入大张量 #

如果你不想换API，可以用tf.train.slice_input_producer先把大张量切成单个样本入队，再用shuffle_batch组合成批次，这样也不会把整个张量塞进图里。

示例代码：

# 先创建切片生产者，每次从大张量中取出一个样本
image_slice, label_slice = tf.train.slice_input_producer(
    [images["train"], labels["train"]],
    shuffle=True,  # 这里先做单样本的shuffle
    seed=self.seed,
    num_epochs=None  # 按需设置训练轮数，None表示无限循环
)

# 再组合成批次
x_train, y_train = tf.train.shuffle_batch(
    [image_slice, label_slice],
    batch_size=self.batch_size,
    capacity=50000,
    min_after_dequeue=10000,  # 建议设为capacity的1/5到1/2，保证shuffle效果，别设0
    num_threads=16,
    allow_smaller_final_batch=True
)

注意这里要把enqueue_many=True去掉，因为slice_input_producer已经是单样本入队了，shuffle_batch默认处理单样本的队列。另外min_after_dequeue设为0的话几乎没有shuffle效果，建议调整成合理值。

额外小提示 #

• 你的num_threads=16，要确保你的CPU核心数能支撑这么多线程，不然可能会出现线程阻塞的情况。
• 如果你的数据集未来还要扩容，建议直接转成TFRecord格式存储，然后用tf.data或者队列读取TFRecord文件，这是处理超大数据集的标准做法，能进一步节省内存。

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