基于TensorFlow学习时OpenCV预处理阶段内存暴涨原因探究

使用TensorFlow训练图像模型时，通过OpenCV读取图像并归一化到0-1区间，过程中内存大幅增长，定位到image = image.astype('float') / 255这行代码占用大量内存，询问该现象的原因。

相关代码：

def create_dataset(img_folder):
    img_data_array=[]
    class_name=[]

    for path in os.listdir(img_folder):
        for file in os.listdir(os.path.join(img_folder, path)):
                continue

            image_path = os.path.join(img_folder, path,  file)

            image = cv2.imread( image_path, cv2.COLOR_BGR2GRAY)

            image = cv2.resize(image, (HEIGHT, WIDTH),interpolation = cv2.INTER_AREA)

            image = np.array(image)
            
            image = image.astype('float') / 255 # <-- 问题代码
         
            img_data_array.append(image)
            class_name.append(path)

    return img_data_array, class_name

• 数据类型扩容导致内存暴增：OpenCV读取的灰度图默认是uint8类型（每个像素占1字节），而astype('float')默认会转成float64（每个像素占8字节），单张图像的内存直接变为原来的8倍。加上除法运算会临时生成一个新的float数组，这会进一步占用额外内存，叠加起来就会导致内存突增。
• 列表累积存储放大内存占用：所有处理后的图像都被append到img_data_array列表中，相当于把所有图像的float64数据全部存在内存里，图像数量越多，内存占用增长越明显。
• 冗余转换增加不必要内存开销：cv2.imread和cv2.resize返回的本身就是numpy数组，没必要再执行image = np.array(image)，这行代码会创建原数组的副本，额外消耗内存。

• 改用更小的浮点类型：把float改成float32，每个像素占4字节，内存占用直接减半：

  image = image.astype('float32') / 255
  

• 移除冗余转换：删掉image = np.array(image)这行代码，避免不必要的数组拷贝。
• 用流式加载替代批量存储：不要用列表一次性存所有图像，改用TensorFlow的tf.data.Dataset实现按需加载，比如：

  import tensorflow as tf
  
  dataset = tf.keras.utils.image_dataset_from_directory(
      img_folder,
      image_size=(HEIGHT, WIDTH),
      color_mode="grayscale",
      batch_size=32,
      label_mode="categorical"
  )
  # 归一化可以放到Dataset的预处理步骤里
  dataset = dataset.map(lambda x, y: (x / 255.0, y))
  

  这种方式不会一次性把所有图像加载到内存，内存占用会大幅降低。
• 尝试原地类型转换：如果原数组后续不再使用，可以用copy=False尝试原地转换，减少临时数组的内存占用：

  image = image.astype('float32', copy=False) / 255
  

内容的提问来源于stack exchange，提问作者김주환