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共七种方式。

（1）统一的请求接口设计

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

核心接口代码如下：

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

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

from transformers import AutoModelForMaskedLMmodel = AutoModelForMaskedLM.from_pretrained('bert-base-uncased')

（2）不同的inference方式

    （2.1）native

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

import torchfrom transformers import AutoModelForMaskedLMmodel = 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用户使用最多的方式。

    （2.2）onnx

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

import torch.onnxdummy_tensor = torch.randint(0, 30522, (1, 512))batch_size = 1torch_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 onnxruntimeonnx_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])

    （2.3）jit

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

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)      

  （2.4）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++编译的各种依赖问题。

  （2.5）其他三种方式暂未测试

（3）评测结果

加载方式	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，需要反复的迭代和优化。

参考资料：

(1)https://github.com/LeeJuly30/BERTCpp

转换pytorch的预训练模型到pb文件，用C++去加载

(2)知乎的BERT加速工作：https://github.com/zhihu/cuBERT

上述工作间接参考本工作，工作比较底层（C++/CUDA）

(3)《那一年，让我整个人升华的C++BERT项目》

讲述了一个小姐姐用BERT做C++改造的故事：

https://www.sohu.com/a/451088664_115128

(4)https://github.com/renatoviolin/BERT-cpp-inference

(5)《Pytorch的C++前端和模型部署》

https://zhpmatrix./2019/03/01/c++-with-pytorch/，笔者很久之前的博客

(6)《直观认识torch.jit模块》

https://zhpmatrix./2019/03/09/torch-jit-pytorch/，笔者很久之前的博客

(7)https://github.com/ShannonAI/service-streamer

工程色彩比较浓重的工作

(8)《模型热更新小记》，与本篇主题不直接相关，但是很有意思

https://mp.weixin.qq.com/s?__biz=MjM5ODkzMzMwMQ==&mid=2650422670&idx=4&sn=da643a8b2810865ab30d2f530077ff06&chksm=becdbbd489ba32c2633e41212a8cae5d48f61f9022705e37997a165583c07f9cd1d5fd912718&mpshare=1&scene=23&srcid=0517NRr5yB17vZnPg65pCl4Y&sharer_sharetime=1623916906151&sharer_shareid=0e8353dcb5f53b85da8e0afe73a0021b%23rd

(9)微软自家的demo

https://github.com/microsoft/onnxruntime/blob/master/onnxruntime/python/tools/transformers/notebooks/PyTorch_Bert-Squad_OnnxRuntime_CPU.ipynb

在该工作中，提供了一些性能测试工具，包括针对ort本身开发的profiler

(10)一个更详细的demo

https://bbs./blogs/180532

(11)快手异构计算团队基于nvidia的fastertransformer进一步做了底层优化，具体包括：算法融合和重构，混合精度量化，内存管理优化，输入padding移除，GEMM配置优化

https://mp.weixin.qq.com/s?__biz=MzA3MzI4MjgzMw==&mid=2650808102&idx=3&sn=cd8f1de3d64c0bfe0e0ba28a526629df&chksm=84e5db58b392524eafed49dd893d064202faebad4cf80beacde641a25a1616d1e83a51354c26&mpshare=1&scene=23&srcid=0201uAERHZKdbPqvF8ACGrPE&sharer_sharetime=1623921005073&sharer_shareid=0e8353dcb5f53b85da8e0afe73a0021b%23rd

(12)微信AI的工作，同样基于nvidia的fastertransformer进一步做底层优化，具体包括：kernel优化(softmax+layernorm)，针对可变序列长度的内存分配算法，针对可变序列长度的batch调度器。

https://github.com/Tencent/TurboTransformers

和快手的工作相比，这篇工作更偏向于工程侧服务过程的优化。不过，二者都很强调对可变长度输入的处理。