使用spaCy en_core_web_trf的nlp.pipe推理时疑似内存泄漏致CUDA溢出

问题背景 #

在16GB显存GPU上运行spaCy nlp.pipe接口的批量推理任务，全程仅执行推理操作，无模型训练、梯度计算相关逻辑。先后调整多组batch_size参数，仍无法避免显存随样本处理量持续上涨的情况，从特征判断疑似存在显存泄漏，最终任务处理到第31000条样本时触发OOM崩溃。
报错详情如下：

CUDA out of memory. Tried to allocate 46.00 MiB (GPU 0; 11.17 GiB total capacity; 10.44 GiB already allocated; 832.00 KiB free; 10.72 GiB reserved in total by PyTorch)

已确认的运行环境 #

• spaCy版本：3.0.5
• 系统平台：Linux-4.19.112+-x86_64-with-Ubuntu-18.04-bionic
• Python版本：3.7.10
• 加载模型：en_core_web_trf 3.0.0

复现代码 #

!python -m spacy download en_core_web_trf

import en_core_web_trf
nlp = en_core_web_trf.load()

# data 为包含10万条句子的数据集数组
data = dataload()

for index, review in enumerate(nlp.pipe(data, batch_size=100)):
    # 执行样本后处理逻辑
    if index % 1000:
        print(index)

可直接落地的解决办法 #

• 强制关闭梯度计算
  这个版本的spaCy transformer组件默认不会关闭PyTorch的梯度计算开关，推理场景下保留计算图完全是多余逻辑，会导致每批样本的计算图不被释放，显存越堆越高。加载模型后先加一行代码关掉梯度开关即可：

  import torch
  torch.set_grad_enabled(False)
  

• 定期清理显存缓存
  每处理固定数量的样本后，主动触发垃圾回收和PyTorch显存清理，避免PyTorch为了预分配缓存机制持续占用显存不释放：

  import gc
  import torch
  
  for index, review in enumerate(nlp.pipe(data, batch_size=32)):
      # 自定义后处理逻辑
      if index % 500 == 0:
          gc.collect()
          torch.cuda.empty_cache()
  

• 禁用不需要的管道组件
  en_core_web_trf默认会加载词性标注、依存分析、命名实体识别、词形还原等全套预测管道，要是推理时只需要用其中部分功能，把不需要的组件关掉能直接大幅降低显存占用：

  # 示例：如果只需要做句法依赖解析，就关掉命名实体、词形还原、属性规则组件
  nlp = en_core_web_trf.load(disable=["ner", "attribute_ruler", "lemmatizer"])
  

• 手动限制transformer组件的真实batch大小
  3.0.x版本的spaCy存在参数传递bug，传给nlp.pipe的batch_size值不会同步到底层transformer组件，组件会自动根据样本长度凑批，遇到长文本很容易一次性凑出超大batch撑爆显存。可以在调用管道的时候通过component_cfg参数给transformer组件单独设置batch大小：

  for index, review in enumerate(
      nlp.pipe(
          data, 
          batch_size=32,
          component_cfg={"transformer": {"batch_size": 32}}
      )
  ):
      # 自定义后处理逻辑
  

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