ocr pipeline 运行内存泄露，且一张图片处理需2s左右，解决方案？

问题分析

您在使用 ocr-detection 模型处理图片时，发现单张图片的处理时间约为 2 秒，并且随着图片数量增加，内存占用显著上升（例如：3100 张图片占用 24GB 内存）。这表明可能存在 性能瓶颈 和 内存泄漏 的问题。以下从加速和解决内存泄漏两个方面提供解决方案。

加速方案

1. 批量处理图片

根据知识库中的内容，批量检测图片可以有效减少请求开销。虽然批量检测可能会拉长平均响应时间，但总体效率会更高。以下是实现批量处理的建议：

• 代码调整： 将单张图片逐一处理改为批量提交多张图片。例如，每次提交 100 张图片进行检测。

  from modelscope.pipelines import pipeline
  from modelscope.utils.constant import Tasks
  
  # 初始化 OCR 检测管道
  p = pipeline('ocr-detection', 'damo/cv_resnet18_ocr-detection-line-level_damo', 
             model_revision='v1.0.0', device="cuda:1")
  
  # 批量处理图片
  batch_size = 100
  for i in range(0, 3000, batch_size):
    batch_results = [p('font1.jpg') for _ in range(batch_size)]
    print(f'Processed {i + batch_size}/3000')
  

• 注意事项：

  • 批量大小需要根据显存容量动态调整，避免显存溢出。
  • 如果模型支持直接传入批量数据（如列表或张量），优先使用该方式以减少重复初始化开销。

2. 优化模型推理

通过以下方法进一步提升推理速度： - 混合精度训练：启用 torch.bfloat16 或 torch.float16 数据类型，降低显存占用并加速计算。 - 模型剪枝与量化：对模型进行剪枝或量化操作，减少参数规模，从而提升推理速度。

示例代码（启用混合精度）：

import torch
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

# 启用混合精度
with torch.cuda.amp.autocast():
    p = pipeline('ocr-detection', 'damo/cv_resnet18_ocr-detection-line-level_damo', 
                 model_revision='v1.0.0', device="cuda:1")
    for i in range(3000):
        result = p('font1.jpg')
        print(f'{i + 1}/3000')

3. 分布式推理

如果单张 GPU 显存不足，可以考虑使用多 GPU 分布式推理。将图片分发到多个 GPU 上并行处理，显著提升吞吐量。

解决内存泄漏方案

1. 手动释放显存

在循环中显式调用 torch.cuda.empty_cache() 清理未使用的显存，避免显存累积占用。

示例代码：

import torch
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

# 初始化 OCR 检测管道
p = pipeline('ocr-detection', 'damo/cv_resnet18_ocr-detection-line-level_damo', 
             model_revision='v1.0.0', device="cuda:1")

for i in range(3000):
    result = p('font1.jpg')
    print(f'{i + 1}/3000')
    
    # 手动清理显存
    if i % 100 == 0:  # 每处理 100 张图片清理一次
        torch.cuda.empty_cache()

2. 检查模型是否正确释放资源

确保每次调用 pipeline 后，模型的中间变量被正确释放。可以通过以下方法排查： - 使用 del 删除不再需要的对象。 - 调用 gc.collect() 强制触发垃圾回收。

示例代码：

import gc
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

# 初始化 OCR 检测管道
p = pipeline('ocr-detection', 'damo/cv_resnet18_ocr-detection-line-level_damo', 
             model_revision='v1.0.0', device="cuda:1")

for i in range(3000):
    result = p('font1.jpg')
    print(f'{i + 1}/3000')
    
    # 删除临时变量并触发垃圾回收
    del result
    gc.collect()

3. 升级 PyTorch 和 CUDA 驱动

根据知识库中的信息，PyTorch 可能因版本不兼容导致内存泄漏问题。建议升级到最新稳定版本的 PyTorch 和 CUDA 驱动程序。

升级命令：

pip install --upgrade torch torchvision torchaudio

同时，确保 GPU 驱动版本与 CUDA 版本匹配。例如： - CUDA 12.1.1 - NVIDIA Driver 535.154.05

总结与建议

1. 加速方案：

   • 使用批量处理减少请求开销。
   • 启用混合精度推理（torch.bfloat16）。
   • 考虑分布式推理以充分利用多 GPU 资源。

2. 内存泄漏解决方案：

   • 定期调用 torch.cuda.empty_cache() 清理显存。
   • 确保模型中间变量被正确释放（del 和 gc.collect()）。
   • 升级 PyTorch 和 CUDA 驱动程序以修复潜在的内存管理问题。

通过上述优化措施，您可以显著提升 OCR 检测的效率并解决内存泄漏问题。