Spacy 推理在处理多个文档时出现 OOM

Question

我正在使用 spacy 处理来自 rest api 的文档。更具体地说，我在 NER 上使用基于变压器的模型 en_core_web_trf，在 GPU 上使用运行。这是与 spacy 相关的代码片段 class（它包含在一些基本的烧瓶服务器中，但我认为这里不重要）

class SpacyExtractor():
    def __init__(self):
        spacy.require_gpu()
        self.model = spacy.load('en_core_web_trf',
                                 disable=["tagger", "parser", "attribute_ruler", "lemmatizer"])


    def get_named_entities(self, text: str):
        doc = self.model(text)
        entities = []
        for ent in doc.ents:
            entities.append((ent.text, ent.label_))
        return entities

问题是，每次调用 get_named_entities，分配的 GPU 内存量都会增加。每次大约 2-3 GB（我通过在应用程序处理文档时反复调用 nvidia-smi 来检查这一点）。所以在打了几次电话之后，我得到了 OOM 错误 RuntimeError: CUDA out of memory. Tried to allocate 2.35 GiB (GPU 0; 10.76 GiB total capacity; 5.02 GiB already allocated; 1.18 GiB free; 8.41 GiB reserved in total by PyTorch) 文件一点也不大，每份1-100页的文字。我想我犯了一些错误，但我只是没有看到它。环境：Ubuntu 18.04，Python 3.8，spacy 3.1.3，cuda 9.1，RTX 2080Ti 11GB 内存

编辑：此外，我在处理单个非常长的文档时发现了 OOM 错误，该文档显示为单个长字符串。

Answer 1

The problem is, with each call of get_named_entities, the amount of GPU memory allocated goes up.

您应该 detach 您的数据如 FAQ 中所述：

Don’t accumulate history across your training loop. By default, computations involving variables that require gradients will keep history. This means that you should avoid using such variables in computations which will live beyond your training loops, e.g., when tracking statistics. Instead, you should detach the variable or access its underlying data.

编辑

您也可以使用

with torch.no_grad():
    doc = self.model(text)

EDIT: Also, I found out the OOM error when processing a single really long document, presented as a single long string.

嗯，这是意料之中的事情。

Spacy 推理在处理多个文档时出现 OOM

Spacy inference goes OOM when processing several documents

python

nlp

spacy

pytorch