使用Hugging Face Transformers做图像特征提取时内存溢出崩溃的解决办法

问题概述 #

使用Hugging Face Transformers进行图像特征提取用于相似性搜索，处理约200张图片后内存占用过高导致系统崩溃；单张图片处理时CPU占用正常，但累计处理数量增加后CPU占用持续上升，内存未按预期回收。

内存泄漏根源分析 #

1. PyTorch计算图与张量未及时释放
   extract_features方法中，未禁用梯度计算会生成冗余计算图，且中间张量（如inputs、image_features）未显式删除，导致内存中残留大量无效张量引用。

2. 图像重复处理引发内存残留
   get_features_for_image函数中已对图像执行resize((224,224))，但extract_features内的get_color_image又重复执行resize操作，生成新的图像对象后原对象引用未及时清理。

3. 类方法与属性命名冲突
   FeatureExtractor类中定义了model()和processor()方法，与初始化时的self.model、self.processor属性重名，会覆盖属性值，导致内存中残留无效对象引用。

4. PyTorch与TensorFlow混合使用的内存冲突
   代码同时使用PyTorch模型（AutoModelForZeroShotImageClassification）和TensorFlow的TFCLIPModel，两个框架内存管理机制独立，易出现内存无法协同回收的问题。

具体修复步骤 #

1. 修复类方法与属性命名冲突 #

修改FeatureExtractor中的方法名，避免与属性重名：

class FeatureExtractor:
    def __enter__(self):
        return self
    def __exit__(self, exc_type, exc_value, traceback):
        pass
    
    def __init__(self, processor=None, model=None, tokenizer=None, text_model=None):
        self.processor = processor
        self.model = model
        self.tokenizer = tokenizer
        self.text_model = text_model

    # 重命名方法，避免与属性冲突
    def get_model(self):
        return self.model
    
    def get_processor(self):
        return self.processor

2. 清理PyTorch计算图与张量 #

在extract_features中禁用梯度计算，显式删除中间张量并清理缓存：

import torch  # 需导入torch库

def extract_features(self, img, grey=False):
    try:
        if grey:
            img = get_greyscale_image(img)
        else:
            img = get_color_image(img)
        inputs = self.processor(images=img, return_tensors="pt")
        with torch.no_grad():  # 禁用梯度计算，避免生成冗余计算图
            image_features = self.model.get_image_features(**inputs)
        image_features /= image_features.norm(dim=-1, keepdim=True)
        feature_np = image_features.detach().numpy()
        # 显式删除中间张量
        del inputs, image_features
        torch.cuda.empty_cache()  # GPU环境下使用，CPU环境可忽略
        return feature_np
    except Exception as e:
        print(e)

3. 移除重复的图像处理操作 #

删除get_features_for_image中重复的resize步骤，避免生成冗余图像对象：

def get_features_for_image(image_meta):
    id = image_meta["id"]
    image_url = image_meta["image_url"]
    try:
        image = get_pil_image_from_url(image_url)
        # 移除重复的resize，extract_features内已处理图像尺寸
        if not os.path.exists(f"features/{id}.npy"):
            image_features = fe.extract_features(image)
            np.save(f"features/{id}.npy", image_features)
            # 显式清理对象
            del image_features, image
            gc.collect()
        return True
    except Exception as e:
        print(f"Error extracting features for image {id}: {e}")

4. 统一深度学习框架，避免混合使用 #

选择单一框架（如全用PyTorch），避免跨框架内存管理冲突：

# 替换TensorFlow模型为PyTorch版本
from transformers import CLIPModel, CLIPTokenizer

# 初始化时统一使用PyTorch模型
model = AutoModelForZeroShotImageClassification.from_pretrained("你的模型名称")
text_model = CLIPModel.from_pretrained("你的模型名称")
tokenizer = CLIPTokenizer.from_pretrained("你的模型名称")

5. 优化上下文管理器（可选） #

在FeatureExtractor的__exit__方法中添加资源清理逻辑，确保退出上下文时释放模型资源：

def __exit__(self, exc_type, exc_value, traceback):
    del self.model, self.processor, self.tokenizer, self.text_model
    torch.cuda.empty_cache()
    gc.collect()

额外优化建议 #

• 批量处理图片：每次处理10-20张图片为一批，处理完成后统一清理内存，减少频繁GC的开销。
• 使用内存监控工具：通过memory_profiler库定位具体代码行的内存占用情况，精准排查泄漏点。
• 压缩特征存储：使用np.savez_compressed替代np.save，减少磁盘占用的同时降低内存中临时数组的大小。

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