TensorFlow创建Dataset时GPU_0_bfc显存不足问题求解

你使用tf.data.Dataset.from_tensor_slices直接传入内存中的numpy数组时，TensorFlow 2.x默认会将传入的全量数组直接拷贝到GPU显存中，这个过程发生在数据集创建阶段，和后续的batch size无关。你的训练集全量大小为5GB，加上CUDA运行时、框架本身的显存开销，超出了你的GPU显存容量，因此会在创建数据集阶段直接触发显存不足报错，调整batch size、修改分配器策略都无法解决这个阶段的问题。

方案1：强制数据集创建在CPU侧（最简便，无需修改现有数据存储逻辑） #

只需要给数据集创建逻辑加上CPU设备上下文，让全量数据保留在内存中，训练时仅按需将单个batch加载到GPU即可，代码修改如下：

import tensorflow as tf
import numpy as np

# 可选：先开启GPU显存按需分配，避免框架预占全量显存
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError as e:
        print(e)

# 强制数据集创建逻辑在CPU执行，全量数据保留在内存
with tf.device('/CPU:0'):
    dataset_train = tf.data.Dataset.from_tensor_slices((X_train, Y_train))
# 批处理后设置预取，自动将下一个batch提前加载到GPU
dataset_train = dataset_train.batch(32).prefetch(tf.data.AUTOTUNE)

# 测试集按相同逻辑处理
with tf.device('/CPU:0'):
    dataset_test = tf.data.Dataset.from_tensor_slices(X_test)
dataset_test = dataset_test.batch(32).prefetch(tf.data.AUTOTUNE)

这个方案可以直接解决你的问题，无需修改现有数据加载逻辑。

方案2：懒加载模式（适合数据集远大于内存/显存的场景） #

如果你的后续数据量还会增长，建议不要提前将全量numpy数组加载到内存，而是将数据拆分为小文件存在硬盘，用生成器方式懒加载，全程不会加载全量数据：

1. 先将全量X_train、Y_train拆分为多个小的.npy文件存到本地，例如每个文件存10000条样本
2. 用生成器构建数据集：

def train_data_generator():
    # 按你拆分的分片数量调整range参数
    for part_idx in range(26):
        x_part = np.load(f"./train_parts/x_train_part_{part_idx}.npy")
        y_part = np.load(f"./train_parts/y_train_part_{part_idx}.npy")
        for x, y in zip(x_part, y_part):
            yield x, y

dataset_train = tf.data.Dataset.from_generator(
    generator=train_data_generator,
    output_signature=(
        tf.TensorSpec(shape=(30,30,3), dtype=tf.float32),
        tf.TensorSpec(shape=(1,), dtype=tf.float32)
    )
).batch(32).prefetch(tf.data.AUTOTUNE)

内容的提问来源于stack exchange，提问作者anastasia-enot