Python读取大量图像并转换为H5文件时如何避免内存占满

核心内存泄漏点排查 #

• 多进程池未正确释放：augment_clean_data_master 函数每次执行都会新建8进程池，用完后未关闭、回收进程资源，进程持有的图像内存会持续残留
• 大变量未完全清理：循环中仅删除了X、y，存储了5万组图像数据的partial_augmented_image_data_list没有被主动释放
• 逻辑冗余重复计算：最后一个分块会重复调用2次augment_clean_data_master，产生双倍的临时内存占用
• numpy/图像处理底层缓存未清理：仅调用gc无法回收部分底层库申请的缓存内存

可直接落地的修复步骤 #

1. 修复多进程池泄漏 #

改写augment_clean_data_master函数，用上下文管理器管理进程池，用完自动回收：

def augment_clean_data_master(master_data):
    import gc
    from multiprocessing import Pool
    data_list = master_data
    print("Length of Data List: ",len(data_list))
    num_processors = 8
    print("Doing Flipped Image Data.")
    # 用with上下文管理器自动关闭进程池
    with Pool(processes = num_processors) as p:
        flipped_output = p_map(get_data_from_line_master_flipped,data_list)
        print("Doing OG Image Data.")
        clean_output = p_map(get_data_from_line_master,data_list)
    print("Merging Data.")
    aug = clean_output + flipped_output
    # 主动删除临时列表再返回
    del flipped_output, clean_output
    gc.collect()
    return aug

2. 修复循环逻辑与变量清理 #

改写循环部分，删除冗余调用，清理所有大变量：

def augment_data_and_save_as_hdf5():
    import gc
    import numpy as np
    import h5py
    import process_data_helpers
    # 提前计算分块总数，避免重复计算
    chunk_size = 25000
    file_name = "./master/master.txt"
    with open(file_name) as f:
        data_list = [x.split(',') for x in f.readlines()]
    # 向上取整计算分块数，自动处理边界
    total_chunks = (len(data_list) + chunk_size - 1) // chunk_size 

    for i in range(total_chunks):
        start_index = i * chunk_size
        # 自动处理最后一块的边界，不用重复判断调用
        stop_index = min((i+1)*chunk_size, len(data_list))
        print("*"*30)
        print(f"Processing chunk {i}, range {start_index}:{stop_index}")
        
        partial_augmented_image_data_list = process_data_helpers.augment_clean_data_master(data_list[start_index:stop_index])

        X = []
        y = []
        for element in partial_augmented_image_data_list:
            try:
                if element[1] != 'None':
                    X.append(element[0])
                    y.append(float(element[1]))
            except:
                pass
        
        X = np.array(X).reshape(-1, 66, 200, 3)
        y = np.array(y).reshape(len(y), 1)
        print("X Shape: ", X.shape)
        print("y Shape: ", y.shape)

        with h5py.File(f'./masterArrays/Data_Chunk_{i}.h5', 'w') as hf:
            hf.create_dataset("X", data=X, compression="gzip", chunks=True, maxshape=(50000,66, 200, 3))
            hf.create_dataset("y", data=y, compression="gzip", chunks=True, maxshape=(50000,1))
        
        # 清理所有当前循环产生的大变量
        del X, y, partial_augmented_image_data_list
        # 强制回收numpy底层缓存
        gc.collect()
        # 如果用了tensorflow/keras，加下面这句清理计算图缓存
        # tf.keras.backend.clear_session()

    print("Augmented & Saved Arrays.")

3. 可选优化降低内存峰值 #

• 把分块大小chunk_size从25000调整为10000~15000，进一步降低单次处理的内存压力
• 把p_map换成imap分批获取进程池返回结果，避免一次性把所有处理好的图像加载到内存中

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