我使用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错误,该文档以单个长字符串的形式呈现
您应该
detach
您的数据,如FAQ中所述:编辑
你也可以使用
这是意料之中的
相关问题 更多 >
编程相关推荐