如何修复TensorFlow Datasets shuffle操作时的内存泄漏问题

基于TensorFlow Datasets下载的Stanford Dog Breed数据集训练模型时，在搭载GPU的Google Colab环境启动训练触发内存错误，导致Colab运行时重启，相关日志输出如下：

tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc:39] Overriding allow_growth setting because the TF_FORCE_GPU_ALLOW_GROWTH environment variable is set. Original config value was 0.

参考TensorFlow官方示例教程编写代码，确认操作执行顺序无误。经定位，数据集shuffle操作是引发内存溢出的根源，且问题仅在调用model.fit()时触发，需要在保留正常shuffle功能的前提下规避内存报错。
问题复现代码如下：

import tensorflow_datasets as tfds

# 加载训练集、测试集划分
(ds_train, ds_test), ds_info = tfds.load('stanford_dogs',
    split=['train', 'test'], shuffle_files=True, as_supervised=True, with_info=True,
)

def normalize_img(image, label):
    return tf.cast(image, tf.float32) / 255., label

ds_train = ds_train.map(normalize_img, num_parallel_calls=tf.data.AUTOTUNE)
# ds_train = ds_train.shuffle(ds_info.splits['train'].num_examples) # 该行触发OOM错误
ds_train = ds_train.batch(batch_size)
ds_train = ds_train.prefetch(tf.data.AUTOTUNE)

ds_test = ds_test.map(normalize_img, num_parallel_calls=tf.data.AUTOTUNE)
ds_test = ds_test.batch(1)
ds_test = ds_test.prefetch(tf.data.AUTOTUNE)

日志中关于GPU allow_growth的内容只是普通提示信息，不是报错根源。
触发OOM的核心原因是shuffle操作传入的缓冲区大小等于训练集全量样本数：Stanford Dogs训练集共12000张图像，shuffle()会将缓冲区大小对应的所有样本提前加载到内存中完成随机打乱，全量加载归一化后的3通道浮点型图像会直接占满Colab运行时的内存上限，触发重启。

• 调小shuffle缓冲区大小，无需将全量样本放入shuffle缓冲区。常规训练场景下，缓冲区大小设置为1000即可满足随机性要求，不会对模型效果产生可感知的影响，内存占用可降至原写法的1/10左右；如果内存仍有压力，可进一步下调至500。
• 调整数据集流水线顺序，将shuffle操作前置到归一化map操作之前。此时shuffle阶段加载的是未解码、未转换为float32格式的原始数据，内存占用远低于预处理后的浮点张量，优化后的流水线顺序参考：

# 先打乱原始样本
ds_train = ds_train.shuffle(1000)
# 再做归一化预处理
ds_train = ds_train.map(normalize_img, num_parallel_calls=tf.data.AUTOTUNE)
ds_train = ds_train.batch(batch_size)
ds_train = ds_train.prefetch(tf.data.AUTOTUNE)

• 如果需要更好的全局打乱效果，无需盲目增大shuffle缓冲区，可在每个训练epoch结束后追加shuffle操作，配合固定随机种子即可同时保证打乱效果和结果可复现。
• 额外优化：测试集不需要执行shuffle操作，当前测试集batch size设为1会大幅拖慢验证速度，可根据显存情况将测试集batch size调整为16或32，不影响验证精度的同时可显著提升运行效率。

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