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

Question

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

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 RAM

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

Answer 1

问题是，每次调用 get_named_entities，分配的 GPU 内存量都会增加。

您应该按照FAQ 中的说明detach 您的数据：

不要在整个训练循环中累积历史记录。默认情况下， 涉及需要梯度的变量的计算将保持 历史。这意味着您应该避免在 超出训练循环的计算，例如，当 跟踪统计。相反，您应该分离变量或访问 其基础数据。

---

编辑

你也可以使用

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

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

嗯，这是意料之中的。