利用 spaCy 和 Cython 提升 NLP 处理速度
自从去年发布了 Python 版的指代消解包(coreference resolution package)以来,许多用户开始使用它来构建各种应用程序,这些应用与最初设计的对话应用有很大不同。
我们发现,尽管在处理对话时该包的速度表现良好,但在处理较大问题时却非常慢。为了找到原因,我启动了一个新项目——NeuralCoref v3.0(https://github.com/huggingface/neuralcoref/)。这个版本的处理速度提高了100倍,每秒可以处理数千个单词,同时保持了准确性和易用性,并且仍然位于 Python 库的生态系统中。
在这篇文章中,我想分享一些在这个项目中所学到的经验,主要包括:
需要注意的是,虽然我们讨论的是 Python,但会涉及到一些 Cython 的技巧。不过,不用担心,Cython 本质上是 Python 的超集,因此并不难掌握!
Python 程序已经可以轻松转变为 Cython 程序。在以下几种情况下,你可能会需要这种加速:
在开始之前,我想指出:这篇文章中的示例代码已上传至 Jupyter Notebook(https://github.com/huggingface/100-times-faster-nlp),你可以尝试一下!
加速第一步:性能分析
首先需要明确的是,大部分纯 Python 代码的表现都是不错的,但有几个瓶颈函数如果能够优化,将会显著提升速度。因此,首先要用分析工具检查 Python 代码,找出慢的地方。一种方法是使用 cProfile:
```python import cProfile import pstats import myslowmodule
cProfile.run('myslowmodule.run()', 'restats') p = pstats.Stats('restats') p.sortstats('cumulative').printstats(30) ```
这可能会揭示出几个循环较慢,特别是当使用神经网络时,可能有几个 NumPy 数组操作较慢。但这里不讨论如何加速 NumPy,已经有大量的文章探讨过这个问题(https://cython.readthedocs.io/en/latest/src/userguide/numpy_tutorial.html)。
接下来,我们该如何加速这些循环?
使用 Cython 加速循环
以一个简单的例子来说明。假设我们有一个包含大量矩形的集合,存储为 Python 对象(即 Rectangle 类的实例)的列表。模块的主要功能是遍历该列表,统计超过某个阈值的矩形数量。
我们的 Python 模块非常简单,如下所示:
```python from random import random
class Rectangle: def init(self, w, h): self.w = w self.h = h
def area(self):
return self.w * self.h
def checkrectangles(rectangles, threshold): nout = 0 for rectangle in rectangles: if rectangle.area() > threshold: nout += 1 return nout
def main(): nrectangles = 10000000 rectangles = [Rectangle(random(), random()) for _ in range(nrectangles)] nout = checkrectangles(rectangles, threshold=0.25) print(n_out) ```
这里的 check_rectangles 函数是瓶颈!它需要遍历大量 Python 对象,而每次循环中 Python 解释器都会执行许多后台工作(如在类中查找 area 方法、打包解包参数、调用 Python API 等),因此这段代码会非常慢。
这里 Cython 可以帮助我们加快循环速度。
Cython 是 Python 的一个超集,它包含两类对象:
高速循环是指 Cython 程序中只访问 Cython C 对象的循环。
设计这种高速循环最直接的方法是定义一个 C 结构,它包含计算过程中所需的所有信息。在这个例子中,该结构需要包含矩形的长和宽。
然后我们可以将矩形列表保存在一个 C 数组中,并传递给 check_rectangle 函数。现在该函数需要接收一个 C 数组作为输入,因此它应该用 cdef 关键字(而不是 def)定义为 Cython 函数。(注意 cdef 也用于定义 Cython C 对象。)
以下是 Cython 高速版本的模块:
```python from cymem.cymem cimport Pool from random import random
cdef struct Rectangle: float w float h
cdef int checkrectangles(Rectangle* rectangles, int nrectangles, float threshold): cdef int nout = 0 # C 数组不包含大小信息 => 我们需要显式给出 for rectangle in rectangles[:nrectangles]: if rectangle.w * rectangle.h > threshold: nout += 1 return nout
def main(): cdef: int nrectangles = 10000000 float threshold = 0.25 Pool mem = Pool() Rectangle* rectangles = mem.alloc(nrectangles, sizeof(Rectangle))
for i in range(n_rectangles):
rectangles[i].w = random()
rectangles[i].h = random()
n_out = check_rectangles(rectangles, n_rectangles, threshold)
print(n_out)
```
这里使用了一个 C 指针数组,但你也可以使用其他方式,如 vectors、pairs、queues 等 C++ 结构(http://cython.readthedocs.io/en/latest/src/userguide/wrapping_CPlusPlus.html#standard-library)。在这段代码中,我还使用了 cymem(https://github.com/explosion/cymem)提供的方便的 Pool() 内存管理对象,这样就不需要手动释放 C 数组了。当 Python 对 Pool 进行垃圾回收时,会自动释放所有通过 Pool 分配的内存。
关于在 NLP 中使用 Cython 的指南,请参阅 spaCy API 的 Cython Conventions(https://spacy.io/api/cython#conventions)。
试一下这段代码
有许多方法可以测试、编译并发布 Cython 代码!Cython 甚至可以像 Python 一样直接用在 Jupyter Notebook 中(http://cython.readthedocs.io/en/latest/src/reference/compilation.html#compiling-notebook)。
首先用 pip install cython 安装 Cython:
在 Jupyter 中测试
在 Jupyter notebook 中通过 %load_ext Cython 加载 Cython 扩展。
现在,只需使用魔术命令(http://cython.readthedocs.io/en/latest/src/reference/compilation.html#compiling-with-a-jupyter-notebook) %%cython 就可以像写 Python 代码一样写 Cython 代码了。
如果在执行 Cython 单元时遇到编译错误,可以在 Jupyter 的终端输出上看到完整的错误信息。
一些常见的错误:如果要编译成 C++(比如使用 spaCy Cython API),需要在 %%cython 后面加入 -+ 标记;如果编译器抱怨 NumPy,需要加入 import numpy 等。
编写、使用并发布 Cython 代码
Cython 代码保存在 .pyx 文件中。这些文件会被 Cython 编译器编译成 C 或 C++ 文件,然后再被系统的 C 编译器编译成字节码。这些字节码可以直接被 Python 解释器使用。
可以在 Python 中使用 pyximport 直接加载 .pyx 文件:
python
import pyximport; pyximport.install()
import my_cython_module
也可以将 Cython 代码构建为 Python 包,并作为正常的 Python 包导入或发布(详细说明见:http://cython.readthedocs.io/en/latest/src/tutorial/cython_tutorial.html#)。
这项工作比较耗时,主要是要处理所有平台上的兼容性问题。如果需要示例的话,spaCy 的安装脚本(https://github.com/explosion/spaCy/blob/master/setup.py)就是一个很好的例子。
在进入 NLP 之前,我们先快速讨论下 def、cdef 和 cpdef 关键字,这些是学习 Cython 时最关键的概念。
Cython 程序中包含三种函数:
def 定义。它的输入和输出都是 Python 对象。内部可以使用 Python 对象,也可以使用 C/C++ 对象,也可以调用 Cython 函数和 Python 函数。cdef 关键字定义。Python 对象和 Cython 对象都可以作为它的输入、输出和内部对象使用。这些函数无法在 Python 空间(即 Python 解释器,和其他需要导入 Cython 模块的纯 Python 模块)中直接访问,但可以被其他 Cython 模块导入。cpdef 定义的 Cython 函数,类似于用 cdef 定义的 Cython 函数,但它们还提供了 Python 封装,因此可以直接在 Python 空间中调用(用 Python 对象作为输入和输出),也可以在其他 Cython 模块中调用(用 C/C++ 或 Python 对象作为输入)。cdef 关键字还有个用法,就是在代码中给 Cython C/C++ 对象定义类型。没有用 cdef 定义类型的对象会被当做 Python 对象处理(因此会降低访问速度)。
通过 spaCy 使用 Cython 加速 NLP
前面说的这些都很好……但这跟 NLP 还没关系呢!没有字符串操作,没有 Unicode 编码,自然语言处理中的难点都没有支持啊!
而且 Cython 的官方文档甚至还反对使用 C 语言级别的字符串(http://cython.readthedocs.io/en/latest/src/tutorial/strings.html):
这就轮到 spaCy 出场了。
spaCy 解决这个问题的办法特别聪明。
将所有字符串转换成 64 比特 hash
在 spaCy 中,所有 Unicode 字符串(token 的文本、token 的小写形式、lemma 形式、词性标注、依存关系树的标签、命名实体标签等)都保存在名为 StringStore 的单一数据结构中,字符串的索引是 64 比特 hash,也就是 C 语言层次上的 uint64_t。
StringStore 对象实现了在 Python unicode 字符串和 64 比特 hash 之间的查找操作。
StringStore 可以从 spaCy 中的任何地方、任何对象中访问,例如可以通过 nlp.vocab.string、doc.vocab.strings 或 span.doc.vocab.string 等。
当模块需要在某些 token 上进行快速处理时,它只会使用 C 语言层次上的 64 比特 hash,而不是使用原始字符串。调用 StringStore 的查找表就会返回与该 hash 关联的 Python unicode 字符串。
但是 spaCy 还做了更多的事情,我们可以通过它访问完整的 C 语言层次上的文档和词汇表结构,因此可以使用 Cython 循环,不需要再自己构建数据结构。
spaCy 的内部数据结构
spaCy 文档的主要数据结构是 Doc 对象,它拥有被处理字符串的 token 序列(称为 words)及所有注解,这些被保存在一个 C 语言对象 doc.c 中,该对象是一个 TokenC 结构的数组。
TokenC(https://github.com/explosion/spaCy/blob/master/spacy/structs.pxd)结构包含关于 token 的所有必要信息。这些信息都保存为 64 比特 hash 的形式,可以通过上面的方法重新构成 unicode 字符串。
看看 spaCy 的 Cython API 文档,就知道这些 C 结构的好处在哪里了。
我们通过一个简单例子看看它在 NLP 处理中的实际应用。
通过 spaCy 和 Cython 进行快速 NLP 处理
假设我们有一个文本文档的数据集需要分析。
```python import urllib.request import spacy
with urllib.request.urlopen('https://raw.githubusercontent.com/pytorch/examples/master/wordlanguagemodel/data/wikitext-2/valid.txt') as response: text = response.read() nlp = spacy.load('en') doc_list = [nlp(text[:800000].decode('utf8')) for _ in range(10)] ```
上面的脚本建立了一个由 10 个 spaCy 解析过的文档组成的列表,每个文档大约有 17 万个词。也可以使用 17 万个文档,每个文档有 10 个词(比如对话的数据集),但那样创建速度就会慢很多,所以还是继续使用 10 个文档好了。
我们要在这个数据集上做一些 NLP 的处理。比如,我们需要计算“run”这个词在数据集中作为名词出现的次数(即被 spaCy 的词性分析(Part-Of-Speech)标记为“NN”的词)。
Python 循环的写法很直接:
```python def slowloop(doclist, word, tag): nout = 0 for doc in doclist: for tok in doc: if tok.lower_ == word and tok.tag_ == tag: nout += 1 return nout
def mainnlpslow(doclist): nout = slowloop(doclist, 'run', 'NN') print(n_out) ```
这也非常慢!在我的笔记本上这段代码大概需要 1.4 秒才能得到结果。如果有 100 万个文档,那就要超过一天的时间。
我们可以使用多任务处理,但在 Python 中通常并不是个好主意(https://youtu.be/yJR3qCUB27I?t=19m29s),因为你得处理 GIL(全局解释器锁,https://wiki.python.org/moin/GlobalInterpreterLock)!
而且,别忘了 Cython 也支持多线程(https://cython.readthedocs.io/en/latest/src/userguide/parallelism.html)!实际上多线程才是 Cython 最精彩的部分,因为 GIL 锁已经被释放,代码可以全速运行了。基本上,Cython 会直接调用 OpenMP。这里不会介绍并行,更多的细节可以参考这里(https://cython.readthedocs.io/en/latest/src/userguide/parallelism.html)。
现在试着用 spaCy 和一点 Cython 加速 Python 代码吧。
首先需要考虑下数据结构。我们需要个 C 层次的数组来保存数据集,其中的指针指向每个文档的 TokenC 数组。还需要将测试字符串(“run”和“NN”)转换成 64 比特 hash。
下面是用 spaCy 编写的 Cython 代码:
```python %%cython -+ import numpy # Sometime we have a fail to import numpy compilation error if we don’t import numpy from cymem.cymem cimport Pool from spacy.tokens.doc cimport Doc from spacy.typedefs cimport hash_t from spacy.structs cimport TokenC
cdef struct DocElement: TokenC* c int length
cdef int fastloop(DocElement* docs, int ndocs, hasht word, hasht tag): cdef int nout = 0 for doc in docs[:ndocs]: for c in doc.c[:doc.length]: if c.lex.lower == word and c.tag == tag: nout += 1 return nout
def mainnlpfast(doclist): cdef int i, nout, ndocs = len(doclist) cdef Pool mem = Pool() cdef DocElement* docs = mem.alloc(n_docs, sizeof(DocElement)) cdef Doc doc
for i, doc in enumerate(doc_list): # Populate our database structure
docs[i].c = doc.c
docs[i].length = doc.length
word_hash = doc.vocab.strings.add('run')
tag_hash = doc.vocab.strings.add('NN')
n_out = fast_loop(docs, n_docs, word_hash, tag_hash)
print(n_out)
```
这段代码有点长,因为得在调用 Cython 函数之前,在 main_nlp_fast 中定义并填充 C 结构。(注:如果在代码中多次使用低级结构,就不要每次填充 C 结构,而是设计一段 Python 代码,利用 Cython 扩展类型(http://cython.readthedocs.io/en/latest/src/userguide/extension_types.html)来封装 C 语言的低级结构。spaCy 的绝大部分数据结构都是这么做的,能优雅地结合速度、低内存占用,以及与外部 Python 库和函数的接口的简单性。)
但它也快得多!在我的 Jupyter notebook 上,这段 Cython 代码只需要大约 20 毫秒,比纯 Python 循环快大约 80 倍。
要知道它只是 Jupyter notebook 单元中的一个模块,还能给其他 Python 模块和函数提供原生的接口,考虑到这一点,它的绝对速度也相当出色:20 毫秒内扫描 1700 万词,意味着每秒能扫描八千万词。
这就是在 NLP 中使用 Cython 的方法,希望你能喜欢。
相关资料
英文原文链接:https://medium.com/huggingface/100-times-faster-natural-language-processing-in-python-ee32033bdced
作者:Thomas Wolf,Huggingface 的机器学习、自然语言处理和深度学习科学负责人。