(转)BERT加速的N种方法

从BERT面世的第二天，笔者就实现了BERT用于序列标注的工作，几乎是全网最早的用BERT做序列标注的工作，到今天离线场景下，BERT做序列标注已经成为一种普惠技术。从huggingface开源Transformers的几乎最早的时间开始跟进，复现组内早期基于Tensorflow做中文纠错的工作，之后模型侧的工作基本一直基于该框架完成。从BERT早期的一系列比较fancy的工作一直在跟进，到组内推广Transformers的使用，到如今PyTorch地位飙升，Transformers社区受众极广，BERT几乎是笔者过去很长一段时间经常讨论的话题。

但是，围绕BERT，最为诟病的一个问题：模型太重，inference时间太长，效果好，但是在线场景基本不能使用？

围绕该问题，学术界和工业界有太多的工作在做。这篇文章简单梳理一些具体的研究方向，同时围绕笔者个人比较感兴趣的一个方向，做一些评测和对比。

那么，具有有哪些研究方向呢？整体上，有两种观察视角。一种是train和inference，另一种是算法侧和工程侧，这里不做具体的区分。

模型大，是慢的一个重要原因，那就换小模型
模型大，通过模型设计，有些部分是可以快的
模型蒸馏
模型压缩剪枝
模型量化：混合精度
服务优化：CPU或者GPU推断，请求管理（批式或者流式），缓存
其他

每个方向都有大量的工作出现，这篇文章主要讨论偏向于工程侧的优化方式。基于huggingface的Transformers的实现，支持不同的模型加载方式native，onnx，jit，libtorch(c++)，native c++(fastertransformer和其他c++版实现)，tensorRT，tensorflow serving共七种方式。

统一的请求接口设计

为了测试不同的inference速度，并不限于模型类型，这里固定模型条件，统一为MaskedLM（bert-base-uncased）。假设脱离本文的主题设定，模型类型显然是影响inference速度的关键因素，这里分为两种条件，第一是不同的模型类型，比如TextCNN和BERT；第二是BERT的不同实现，比如Layer数量的不同，特殊Trick的使用等。
核心接口代码如下：

#预处理：载入分词器
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
#测试文本
text = "[CLS] In deep [MASK], everything is amazing![SEP]"
#分词
tokenized_text = tokenizer.tokenize(text)
#token2id
indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)
#服务请求
url=""
post(url)

得益于Transformers的优雅的接口设计，可以利用两行代码加载分词器，类似的，可以用两行代码加载模型：

1 2	from transformers import AutoModelForMaskedLM model = AutoModelForMaskedLM.from_pretrained("bert-base-uncased")

不同的inference方式

native

朴素的方式是直接加载pytorch_model.bin，config.json， vocab.txt，作为server端的模型。核心服务代码如下：

import torch
from transformers import AutoModelForMaskedLM
model = AutoModelForMaskedLM.from_pretrained("bert-base-uncased")
with torch.no_grad():
        output = model(tokens.to(device))[0].detach().cpu()
output = torch.softmax(output[0][idx_to_predict], 0)

这种方式是平时pytorch用户使用最多的方式。

onnx

笔者第一次接触onnx是2018年做CV的时候，那个时候需要将一个PyTorch的模型转化为onnx，做android移动端的部署，大概在那个时候，不同框架之间的模型转化已经成为一个业界的实际需求。为了通过onnx加载模型，首先需要将native的模型转化为onnx的模型。模型转换代码如下：

import torch.onnx
dummy_tensor = torch.randint(0, 30522, (1, 512))
batch_size = 1
torch_out = model(dummy_tensor)
torch.onnx.export(model,               # model being run
                  dummy_tensor,        # model input (or a tuple for multiple inputs)
                  model_path,   # where to save the model (can be a file or file-like object)
                  export_params=True,     # store the trained parameter weights inside the model file
                  opset_version=10,          # the ONNX version to export the model to
                  do_constant_folding=True,  # whether to execute constant folding for optimization
                  input_names=['input'],   # the model's input names
                  output_names=['output'],  # the model's output names
                  dynamic_axes={'input': {0: 'batch_size'},    # variable length axes
                                'output': {0: 'batch_size'}})

这里转换的逻辑中，有一个细节。导入模型之后，需要通过构造一个dummy tensor才能够获取网络的结构，同时模型转换中提供了一些优化的方式。

核心服务代码如下：

import onnxruntime
onnx_session = onnxruntime.InferenceSession(model_path)
ort_inputs = {onnx_session.get_inputs()[0].name: tokens}
ort_outs = onnx_session.run(None, ort_inputs)
output = np.array(ort_outs)[0][0]
output = softmax(output[idx_to_predict])

jit

使用jit的方式，同样需要做模型转换，转换代码如下：

1
2
3

with torch.no_grad():
    traced_model = torch.jit.trace(model, dummy_tensor)
    torch.jit.save(traced_model, model_path)

核心服务代码如下：

model = torch.jit.load(model_path)
with torch.no_grad():
        output = model(tokens.to(device))[0].detach().cpu()
 output = torch.softmax(output[0][idx_to_predict], 0)

libtorch(c++)

采用libtorch(c++)加载的模型同jit，服务端的核心加载代码如下：

#include "torch/script.h"
torch::jit::script::Module module = torch::jit::load(model_path)
module.eval()
module.forward(tokens)

这里值得一提的是，不同于python的server端，可以选择fastAPI，flask，gunicorn等，c++也有对应的server端，典型的比如crow。使用该种方式的一个问题是：要解决c++编译的各种依赖问题。

备注：其他三种方式暂未测试

评测结果

加载方式	onnx	native	jit	libtorch(c++)	备注
时间(相同请求次数)	6.20s	7.07s	6.83s	libtorch(c++)，限于各种依赖，笔者未测	笔者的结果
时间(相同请求次数)	12.43s	19.43s	18.24s	12.10s	他人的结果

笔者个人的环境和他人的环境不相同，因此具体时间上不同，但是趋势是基本一致的。onnx和libtorch(c++)的方式都较快，NLP算法同学中，python用户居多，因此选择onnx是一种比较理想的方式。native是最慢的，也就是说最常用的方式恰恰是inference效率最低的方式。jit介于两者之间。

对于没有实测过的结果，这里给出一张他人的评测结果，如下：

对比可知：说啥都没有用C++重写一遍来的快！笔者在之前做过一个表格数据处理的加速，向量指令，cache等多种技术都有尝试，最后发现，C++重写一遍核心逻辑，速度立刻显著提升。关于BERT的C++实现，可以参考字节的开源工作。

说了辣么多，咋整吧？只能具体情况具体分析了。从整体上看BERT的加速可以从多个方面开展。但是围绕这篇文章的主题，C++的加速方式效果最理想，但是成本也较高。onnx的方法目前来看，成本较低可执行。实际上，最近的天池的小布助手比赛中，Top选手也多采用了这种方案，但是采用tensorRT的方式也有，这篇文章没有做实测，可以作为一种备选的方案。此外，配合低精度，服务优化等方式。不论怎样，从一开始，结合对数据的理解，选择一个小的模型，使用最native的方式也许就可以满足inference的要求了。蒸馏和剪枝在技术上比较fancy，需要反复的迭代和优化。

转：The following article is from KBQA沉思录 Author zhpmatrix