在处理多个文档时,空间推理会失效

2024-05-13 02:36:52 发布

您现在位置:Python中文网/ 问答频道 /正文

我使用spacy来处理通过RESTAPI传递的文档。更具体地说,我使用基于transformer的模型en_core_web_trf作为NER,运行在GPU上。下面是spacy相关类的代码片段(它打包在一些基本的flask服务器中,但我认为这并不重要)

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检查了这一点)。所以打了几个电话后,我发现了一个错误 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错误,该文档以单个长字符串的形式呈现


Tags: text文档coreselfwebmodelgpuspacy
1条回答
网友
1楼 · 发布于 2024-05-13 02:36:52

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.

这是意料之中的

相关问题 更多 >