使用tf.data.Dataset.from_tensor_slices时GPU内存异常占用问题咨询

为什么内存需求远大于数据集实际大小？ #

• 数据类型膨胀：原始图像是uint8格式（每个像素1字节），转成float32后每个像素占4字节，加上RGB3通道，单张图像内存从128×128×3=49152字节变成128×128×3×4=196608字节。80000张的话，总内存约14.6GiB；如果误使用float64，内存会直接翻倍到29GiB左右，正好匹配你遇到的错误。
• TensorFlow张量的存储逻辑：用tf.data.Dataset.from_tensor_slices(data_imgs)时，整个numpy数组会被转为TensorFlow张量。如果代码默认在GPU环境运行，TensorFlow会尝试把完整张量加载到GPU内存，而非分批加载，直接占满所需内存。
• 额外内存开销：TensorFlow处理张量时，还会产生索引、临时缓存等额外内存占用，进一步加剧内存压力。

解决办法 #

• 修正数据类型：确保归一化后的数据是float32而非float64，转numpy数组时指定dtype=np.float32，直接把内存占用减半。
• 避免一次性加载全部数据：改用从文件路径按需加载的方式，让TensorFlow分批读取图像，不会一次性占满内存。示例代码：

  file_paths = tf.data.Dataset.list_files("path/to/your/images/*") # 替换为你的图像路径通配符
  def load_and_preprocess(path):
      img = tf.io.read_file(path)
      img = tf.image.decode_png(img, channels=3) # 根据图像格式切换decode函数，比如decode_jpeg
      img = tf.cast(img, tf.float32) / 255.0 # 归一化到[0,1]
      return img
  dataset = file_paths.map(load_and_preprocess, num_parallel_calls=tf.data.AUTOTUNE)
  dataset = dataset.batch(64).prefetch(tf.data.AUTOTUNE)
  

• 启用GPU内存动态增长：在代码开头添加以下配置，让TensorFlow按需分配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)
  

• 强制分批处理：如果必须用已加载的numpy数组创建数据集，立刻添加batch()操作，限制单次加载到内存的数据量：

  dataset = tf.data.Dataset.from_tensor_slices(data_imgs).batch(64)
  

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