使用tf.data.Dataset.from_tensor_slices加载80k图像时GPU报错排查

核心原因 #

你遇到的错误本质是一次性将80k张图像加载为单个张量并尝试向GPU传输时，TensorFlow的张量初始化流程出现异常，并非单纯的显存/内存不足：

1. tf.data.Dataset.from_tensor_slices(data_imgs)会把所有图像数据打包成一个巨型CPU张量，当数据量达到80k时，Eager模式下TensorFlow会尝试自动将这个张量复制到GPU，但巨型张量的初始化过程中容易出现设备间传输的隐性错误（哪怕显存使用率看起来不高）。
2. 即便CPU内存有64GB，巨型张量的创建、shuffle缓冲区的分配也可能引发内存碎片化，间接导致GPU端的张量初始化失败。

具体解决步骤 #

1. 改用按需加载（推荐） #

放弃直接从内存张量创建Dataset，改为从文件路径读取，避免一次性加载所有数据：

# 假设你的图像文件都在目标目录下，示例为"./dataset/"
file_paths = tf.data.Dataset.list_files("./dataset/*.png", shuffle=True)

def load_and_preprocess(path):
    img = tf.io.read_file(path)
    img = tf.image.decode_png(img, channels=3)
    img = tf.image.resize(img, [128, 128])
    img = (img / 255.0) - 0.5  # 按你的实际预处理逻辑调整
    return img

train_dataset = file_paths.map(load_and_preprocess, num_parallel_calls=tf.data.AUTOTUNE)\
                          .shuffle(10000)\
                          .batch(1000)\
                          .prefetch(tf.data.AUTOTUNE)

2. 若必须用内存中的数据 #

如果data_imgs已经在内存中，调整Dataset的创建方式，强制数据留在CPU处理，避免自动向GPU传输：

# 创建Dataset时指定CPU设备
with tf.device('/CPU:0'):
    train_dataset = tf.data.Dataset.from_tensor_slices(data_imgs)\
                                  .shuffle(10000)\
                                  .batch(1000)\
                                  .prefetch(tf.data.AUTOTUNE)

同时可尝试：

• 降低BATCH_SIZE（比如从1000降到500），减少单次向GPU传输的数据量
• 检查data_imgs的张量形状是否完全统一（80k张图中可能混入了尺寸/通道数不一致的图像，导致张量创建异常）

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