SpaCy训练NER时GPU内存利用率极低，请求技术协助

我在训练SpaCy NER模型时遇到GPU利用率极低的问题，训练期间GPU利用率仅11-13%，接近空闲状态。已通过Torch验证服务器GPU可正常分配全部8GB显存，硬件运行正常，问题出在SpaCy代码逻辑中。

我的训练代码：

spacy.require_gpu()
nlp = spacy.blank('en')

if 'ner' not in nlp.pipe_names:
    ner = nlp.add_pipe('ner')
else:
    ner = nlp.get_pipe('ner')

docs = load_data(ANNOTATED_DATA_FILENAME_BIN)
train_data, test_data = split_data(docs, DATA_SPLIT)

unique_labels = set(ent.label_ for doc in train_data for ent in doc.ents)
for label in unique_labels:
    ner.add_label(label)

optimizer = nlp.initialize()

for i in range(EPOCHS):
    print(f"Starting Epoch {i+1}...")
    losses = {}
    batches = minibatch(train_data, size=compounding(4., 4096, 1.001))
    for batch in batches:
        for doc in batch:
            example = Example.from_dict(doc, {'entities': [(ent.start_char, ent.end_char, ent.label_) for ent in doc.ents]})
            nlp.update([example], drop=0.5, losses=losses, sgd=optimizer)
    print(f"Losses at iteration {i}: {losses}")

GPU状态截图：

核心问题分析 #

1. 单样本更新浪费GPU并行能力：代码对批次内的每个文档单独创建Example并调用nlp.update([example])，每次仅传入单个样本。GPU的优势是批量并行处理，这种单样本更新会让GPU大部分时间处于等待数据的状态，直接导致利用率极低。
2. 未利用minibatch的批次特性：minibatch生成的是批量文档集合，但代码没有批量转换为Example并一次性传入更新，完全浪费了批次处理的效率。

修正后的代码 #

spacy.require_gpu()
nlp = spacy.blank('en')

if 'ner' not in nlp.pipe_names:
    ner = nlp.add_pipe('ner')
else:
    ner = nlp.get_pipe('ner')

docs = load_data(ANNOTATED_DATA_FILENAME_BIN)
train_data, test_data = split_data(docs, DATA_SPLIT)

unique_labels = set(ent.label_ for doc in train_data for ent in doc.ents)
for label in unique_labels:
    ner.add_label(label)

optimizer = nlp.initialize()

for i in range(EPOCHS):
    print(f"Starting Epoch {i+1}...")
    losses = {}
    # 调整批次增长范围，避免过大批次导致显存溢出
    batches = minibatch(train_data, size=compounding(4., 64, 1.001))
    for batch in batches:
        # 批量生成Example
        examples = []
        for doc in batch:
            annotations = {'entities': [(ent.start_char, ent.end_char, ent.label_) for ent in doc.ents]}
            examples.append(Example.from_dict(doc, annotations))
        # 一次性传入整个批次的样本进行更新
        nlp.update(examples, drop=0.5, losses=losses, sgd=optimizer)
    print(f"Losses at iteration {i}: {losses}")

额外优化建议 #

• 合理调整批次大小：原代码中最大批次4096可能远超GPU显存承载能力，建议根据实际显存调整（比如最大64或128），平衡显存占用与GPU利用率。
• 验证GPU组件运行状态：初始化后可添加以下代码确认NER组件是否在GPU上运行：

  print(f"NER pipe is using GPU: {ner._model.is_device('cuda')}")
  

• 固定随机种子：添加spacy.util.fix_random_seed(42)保证训练可复现，同时避免随机性导致的效率波动。

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