cv.imread加载54K张512×512 PNG图像时出现内存不足报错如何解决？

首先计算全量加载的内存占用基线：单张512×512×3的uint8格式图像占约768KB，54K张总大小约40GB；对应512×512×2的mask单张约512KB，54K张总大小约27GB，两者合计接近70GB，远高于普通消费级设备的内存容量，直接全量加载OOM是必然结果。从内存管理角度可采用以下优化方案：

• 用批次加载生成器替换全量加载：这是最高效的解决方案，不需要把全量数据存入内存，仅在训练每一步时加载当前批次需要的样本。可以继承tf.keras.utils.Sequence类实现自定义数据加载器，在__getitem__方法中完成单批次图片、mask的读取和预处理，训练时直接将生成器实例传入model.fit()即可。该方案的内存占用仅和批次大小相关，比如批次设为16时，单轮加载的内存占用不到30MB，完全规避全量加载的内存压力。
• 优化现有全量加载逻辑的内存开销：如果确实需要全量加载数据到内存，可以改掉先存list再转numpy数组的写法：提前预分配对应大小的numpy数组，直接将读取的图片写入预分配数组的对应位置，避免list存储所有单张图片产生的双倍内存占用。示例写法：

import numpy as np
import cv2
import glob
import gc

img_files = glob.glob(r'C:\Users\77901677\Projects\images1\*.png')
# 提前预分配数组，不要用list append
dataset_image = np.zeros((len(img_files), 512, 512, 3), dtype=np.uint8)
for idx, path in enumerate(img_files):
    dataset_image[idx] = cv2.imread(path)
    if idx % 1000 == 0:
        gc.collect()

mask_files = glob.glob(r'C:\Users\77901677\Projects\annotations1\*.png')
dataset_mask = np.zeros((len(mask_files), 512, 512, 2), dtype=np.uint8)
for idx, path in enumerate(mask_files):
    mask = cv2.imread(path)
    dataset_mask[idx] = mask[:,:,1:]
    del mask
    if idx % 1000 == 0:
        gc.collect()

每加载1000~2000张调用一次gc.collect()可以手动回收内存碎片，减少无效内存占用。

• 用内存映射文件处理全量数据：如果需要频繁访问全量数据集、不想每次训练都重复读磁盘，可以先把数据集存为npy格式，后续加载时开启内存映射模式：

# 首次运行保存数据集（仅需执行一次）
np.save('dataset_image.npy', dataset_image)
np.save('dataset_mask.npy', dataset_mask)
# 后续训练加载时使用内存映射
dataset_image = np.load('dataset_image.npy', mmap_mode='r')
dataset_mask = np.load('dataset_mask.npy', mmap_mode='r')

内存映射模式下数据实际存储在磁盘上，仅当前访问的分片会被加载到内存，使用体验和全量加载到内存的数组完全一致，内存占用极低。

• 延后数据类型转换：加载阶段始终使用uint8格式存储图像和mask，不要提前做归一化转float32格式，否则内存占用会直接翻4倍。归一化、标准化这类预处理操作可以放到批次加载逻辑里，或者直接在模型的首层加入Rescaling层完成，完全不影响训练效果。

内容的提问来源于stack exchange，提问作者Hector Arroyo Gzlez