Python代码执行内存占用过高问题排查与优化咨询

内存占用过高的原因 #

1. HDFStore缓存与重复加载累积：每次循环读取vector[img]会将对应Series完整加载到内存，且Pandas的HDFStore默认缓存读取数据，随着处理图片数量增加，缓存数据持续累积，占用大量内存。
2. 临时对象未及时回收：循环中生成的特征数组、相似度列表array、Pandas Series等临时对象，未被手动清理，Python垃圾回收无法立刻释放内存，导致占用逐步升高。
3. Pandas Series额外开销：用Series存储单张图片特征，相比原生NumPy数组会引入额外对象开销，处理2000张图时，这类小对象的内存消耗会被放大。
4. 重复读取当前图特征：比对循环中反复从HDFStore读取当前图的512维特征，重复加载数据到内存，加剧内存占用。

代码中的错误/冗余操作 #

1. 冗余的HDF初始化：通过open('vector.h5', 'wb')手动清空文件完全多余，Pandas HDFStore用w模式打开时会自动创建/清空文件，当前r+模式下的手动清空属于无效操作。
2. 重复获取HDF键列表：每次循环调用vector.keys()会遍历HDF所有键，产生额外IO和内存开销，可维护一个内存中的键列表逐步更新，无需每次重新读取。
3. 未复用当前图特征：比对时反复从HDF读取当前图特征，未提前存入内存变量复用，导致重复加载数据。
4. 未显式回收内存：循环中临时对象未手动删除，也未触发垃圾回收，内存无法及时释放。

优化方案 #

1. 优化HDF操作逻辑 #

• 复用当前图特征：处理单张图时，先将其512维特征存入内存变量，避免比对循环中反复读取HDF：

  # 处理当前图后，把特征存到内存
  current_emb = vector[img][:512].values
  # 比对时直接用current_emb
  metrics_vector1 = compare_vector_cosine(current_emb, vector[img_vector][:512].values)
  

• 清理HDF缓存：初始化HDFStore时启用压缩，或定期清理缓存：

  vector = pd.HDFStore('vector.h5', mode='r+', complib='zlib', complevel=5, chunksize=100)
  # 每处理100张图清理一次缓存
  if idx % 100 == 0:
      vector.clear_cache()
      compare.clear_cache()
  

2. 用h5py替代Pandas降低内存开销 #

直接使用底层h5py库操作HDF5，避免Pandas Series的额外内存开销，直接存储和读取NumPy数组：

import h5py
import numpy as np

# 初始化HDF文件（无需手动清空）
with h5py.File('vector.h5', 'w') as f:
    pass
with h5py.File('compare.h5', 'w') as f:
    pass

# 维护内存中的已处理图片列表
processed_imgs = []

for idx, img in enumerate(IMAGE_TITLE):
    # 获取特征
    embedding_main = get_insightface_embedding(PATH_DIR + img)
    analyze = get_deepface_analyze(PATH_DIR + img)
    embedding_second = get_face_recognition_embedding(PATH_DIR + img)
    
    # 拼接为NumPy数组
    full_vec = np.concatenate([
        embedding_main['embedding'],
        embedding_second,
        embedding_main['pose'],
        np.array([embedding_main['gender'], embedding_main['age'], analyze['dominant_race']])
    ], axis=0)
    
    # 写入特征到HDF5
    with h5py.File('vector.h5', 'a') as f:
        f.create_dataset(img, data=full_vec, compression='gzip')
    processed_imgs.append(img)
    
    # 获取当前图的512维特征
    current_emb = full_vec[:512]
    # 预分配相似度数组
    sim_array = np.empty(len(processed_imgs), dtype=np.float32)
    
    # 读取已存特征并计算相似度
    with h5py.File('vector.h5', 'r') as f:
        for i, key in enumerate(processed_imgs):
            vec = f[key][:512]
            sim_array[i] = compare_vector_cosine(current_emb, vec)
    
    # 写入相似度结果
    with h5py.File('compare.h5', 'a') as f:
        f.create_dataset(img, data=sim_array, compression='gzip')
    
    # 手动清理临时对象并触发垃圾回收
    del embedding_main, analyze, embedding_second, full_vec, sim_array, current_emb
    import gc
    gc.collect()

3. 优化比对逻辑与内存管理 #

• 预分配相似度数组：用np.empty提前分配数组空间，替代列表append，减少内存扩容开销。
• 分批处理：每处理100张图就关闭并重新打开HDFStore，强制清理缓存。
• 启用HDF压缩：存储时启用compression='gzip'，减少磁盘占用的同时降低读取内存开销。

4. 避免冗余计算 #

处理第N张图时，只需与前N-1张图比对（之前的图已完成更早图片的比对），仅计算新增相似度并补充到已有记录，减少重复计算和存储量。

原代码 #

with open('vector.h5', 'wb') as f:
    f.close()

vector = pd.HDFStore('vector.h5', mode='r+')


with open('compare.h5', 'wb') as f:
    f.close()

compare = pd.HDFStore('compare.h5', mode='r+')

for idx, img in enumerate(IMAGE_TITLE):
    
    embedding_main = get_insightface_embedding(PATH_DIR + img) #512 numbers type float
    analyze = get_deepface_analyze(PATH_DIR + img) #string length 15
    embedding_second = get_face_recognition_embedding(PATH_DIR + img) #128 numbers type float

    vector[img] = pd.Series(np.concatenate([embedding_main['embedding'], 
                                            embedding_second,
                                            embedding_main['pose'],
                                            np.array([embedding_main['gender'], embedding_main['age'], analyze['dominant_race']])], axis=0))
    imgs_posted = vector.keys()
    
    array = []
    
    for img_vector in imgs_posted:

        metrics_vector1 = compare_vector_cosine(vector[img][:512], vector[img_vector][:512])
        
        array.append(metrics_vector1)
     
    compare[img] = pd.Series(data=np.array(array), index=imgs_posted)

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