Spacy内存占用过高问题排查与性能优化方案咨询

我有数万份文档，每份约含150k字符、25k空格分隔token、2k唯一token，用Spacy做预处理（停用词移除+词形还原），预处理依赖token.pos_和token.lemma_。之前误禁用了POS标注所需的tok2vec组件，修复后内存占用暴增，应用挂起后被OOM killer终止。

当前用分块（每块100份文档）、n_process=4的方式把文档送入nlp.pipe处理，修复前运行正常，现在只有分块数降到25-50才能正常运行，降低n_process到1也没明显改善。相关代码如下：

import spacy
from bs4 import BeautifulSoup
import unidecode
import re

nlp = spacy.load('en_core_web_lg')
nlp.max_length = 5000000
nlp.disable_pipe("parser")
nlp.disable_pipe("ner")
nlp.enable_pipe("senter")

def pre_pre_process(record, synswap=True):
    (doc_id, text) = record

    # partial pre-preprocessing = just strip HTML
    text1 = BeautifulSoup(text, "html.parser").get_text(separator=" ")

    # full pre-preprocessing = do all the pre-preprocessing
    text2 = " ".join(text1.strip().split())
    text2 = unidecode.unidecode(text2)
    text2 = text2.lower()
    
    return (text2, {'doc_id': doc_id, 'strip_html': text1, 'ppp': 'full-ppp'})


def pre_process_text(doc, convert_num=True, lemmatization=True,
                     punctuations=True, remove_num=True, special_chars=True,
                     stop_words=True, short_char=True, remove_edgar_junk=True):
    fully_processed = []
    edgar_jnk_patt = re.compile('(?is)ex-\d+\.?\d*')  # noqa: W605
    edgar_jnk = []

    for token in doc:
        # (token, token.pos_, token.is_stop, token.is_punct, token.lemma_)
        flag = True  # assume every token should be added to the vocab
        edit = token.text
        # remove stop words
        if stop_words is True and token.is_stop and token.pos_ != 'NUM':
            flag = False
        # remove punctuations
        if punctuations is True and (token.pos_ == 'PUNCT' or token.is_punct) and flag is True:
            flag = False
        # remove special characters
        if special_chars is True and token.pos_ == 'SYM' and flag is True:
            flag = False
        # remove numbers
        if remove_num is True and (token.pos_ == 'NUM' or token.text.isnumeric()) and flag is True:
            flag = False
        # remove short tokens
        if short_char is True and len(token) < 3 and flag is True:
            flag = False
        # convert tokens to base form
        elif lemmatization is True and token.lemma_ != "-PRON-" and flag is True:
            edit = token.lemma_
        # remove edgar junk
        if remove_edgar_junk is True:
            if token.i < 10:
                if token.text.endswith(('.htm', '.html')):
                    flag = False
                    edgar_jnk.append(token.lemma)
                elif edgar_jnk_patt.search(token.lemma_):
                    flag = False
                    edgar_jnk.append(token.lemma)
            if token.lemma in edgar_jnk and flag is True:
                flag = False

        # append tokens edited and not removed to list
        if edit != "" and flag is True:
            fully_processed.append(edit)
    return fully_processed

# 完整脚本中，data从数据库查询，每次取query_limit=50条，循环直到无数据可查

completed_jobs = []
pipeline_texts = [pre_pre_process(d) for d in data]
for doc, context in nlp.pipe(pipeline_texts, as_tuples=True, n_process=4):
    tokens = pre_process_text(doc)
    completed_jobs.append((context, tokens))

1. tok2vec组件为何占用大量内存？ #

• tok2vec是Spacy生成token向量的核心组件，en_core_web_lg模型的tok2vec依赖大型预训练词向量表（约1M词汇，每个向量300维），加载时会将整个词表载入内存。
• 启用tok2vec后，每个Doc对象会存储token的向量表示；加上单份文档token数多达25k，分块+多进程处理时，每个子进程会复制一份模型实例，内存开销叠加明显。
• 此前禁用tok2vec时，模型跳过了向量生成和POS标注计算，内存占用自然极低；修复后POS标注依赖tok2vec的向量输出，必须加载完整组件，内存占用直接飙升。

2. 如何分析nlp.pipe中的内存占用情况？ #

• 进程级监控：用ps或top命令实时查看各Spacy子进程的内存占用，重点关注RSS（常驻内存）指标，判断是否为多进程复制模型导致的内存叠加。
• 代码内统计：用tracemalloc模块记录内存变化，在处理循环中定期打印内存使用：

  import tracemalloc
  tracemalloc.start()
  
  completed_jobs = []
  pipeline_texts = (pre_pre_process(d) for d in data)
  for i, (doc, context) in enumerate(nlp.pipe(pipeline_texts, as_tuples=True, n_process=4)):
      tokens = pre_process_text(doc)
      completed_jobs.append((context, tokens))
      if i % 10 == 0:
          current, peak = tracemalloc.get_traced_memory()
          print(f"Processed {i} docs - Current memory: {current/1e6:.2f}MB; Peak: {peak/1e6:.2f}MB")
  
  tracemalloc.stop()
  

• Spacy内置工具：用spacy debug data分析单份文档的处理开销，或spacy validate检查模型组件加载情况，确认是否有冗余组件占用内存。

3. 该预处理流水线是否有更优实现方式？ #

• 优化分块与进程配置：
  • 进一步降低单分块文档数量（如20-30份），避免单进程同时处理过多大文档导致内存溢出。
  • 给nlp.pipe设置batch_size参数（如batch_size=10），让Spacy更精细地控制内存分配。
• 延迟加载与内存清理：
  • 用生成器替代列表存储预处理后的文本，避免一次性加载所有数据到内存：

    pipeline_texts = (pre_pre_process(d) for d in data)  # 生成器而非列表
    for doc, context in nlp.pipe(pipeline_texts, as_tuples=True, n_process=4, batch_size=10):
        tokens = pre_process_text(doc)
        completed_jobs.append((context, tokens))
        del doc  # 手动释放Doc对象，辅助GC回收
    

• 改用轻量模型：若无需大型词向量，替换为en_core_web_sm模型，其tok2vec组件更小，内存占用大幅降低，且不影响POS标注和词形还原功能。

4. 预处理逻辑是否有更优方案（Spacy是否有内置实现，当前方案是否标准？） #

• Spacy内置组件替代部分逻辑：
  • 停用词移除：token.is_stop是Spacy标准实现，当前逻辑合理；可通过nlp.select_pipes(enable=["tok2vec", "tagger", "lemmatizer", "senter"])明确启用所需组件，避免冗余加载。
  • 词形还原：token.lemma_是Spacy标准实现，跳过代词的处理符合常规需求。
• 简化过滤逻辑：
  • 合并条件判断，减少重复的flag检查，提升执行效率：

    for token in doc:
        # 合并过滤条件，直接跳过不符合要求的token
        if (stop_words and token.is_stop and token.pos_ != 'NUM') or \
           (punctuations and (token.pos_ == 'PUNCT' or token.is_punct)) or \
           (special_chars and token.pos_ == 'SYM') or \
           (remove_num and (token.pos_ == 'NUM' or token.text.isnumeric())) or \
           (short_char and len(token) < 3):
            continue
        
        # 词形还原处理
        edit = token.lemma_ if (lemmatization and token.lemma_ != "-PRON-") else token.text
        
        # Edgar垃圾内容过滤
        if remove_edgar_junk:
            if token.i < 10:
                if token.text.endswith(('.htm', '.html')) or edgar_jnk_patt.search(token.lemma_):
                    edgar_jnk.append(token.lemma)
                    continue
            if token.lemma in edgar_jnk:
                continue
        
        if edit:
            fully_processed.append(edit)
    

• 自定义管道组件：将预处理逻辑封装为Spacy自定义管道，与nlp.pipe无缝集成，减少手动循环开销：

  @spacy.Language.component("custom_preprocessor")
  def custom_preprocessor(doc, **kwargs):
      processed_tokens = pre_process_text(doc, **kwargs)
      doc._.processed_tokens = processed_tokens
      return doc
  
  # 添加到管道末尾
  nlp.add_pipe("custom_preprocessor", last=True)
  
  # 使用时直接获取处理结果
  for doc, context in nlp.pipe(pipeline_texts, as_tuples=True, n_process=4):
      completed_jobs.append((context, doc._.processed_tokens))
  

• 当前方案的合理性：你的预处理流程（清洗、过滤、词形还原）是NLP预处理的标准流程，可通过Spacy内置管道进一步优化执行效率和内存管理。

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