基于TensorRT的BERT推断加速与服务部署-技术圈

BERT的出现真是广大NLPer的福音，在很多任务上能取得显著提升。不例外，作者在工作过程中也使用了BERT进行下游任务训练，但在感叹BERT真香的时候，它及其漫长的推断时间让人感到很为难。本文就记录了在使用tensorRT部署BERT时候的各种坑。话不多说，先给下最终模型推断时间对比（如下面表格所示），然后开始我们的填坑记。

max_seq	batch_size	1	5	10	20	40	100
128	Tensorflow(ms)	20	24	45	76	130	300
128	tensorRT(ms)	2.2	5.2	7.2	10.1	19	51

实战系列篇章中主要会分享，解决实际问题时的过程、遇到的问题或者使用的工具等等。如问题分解、bug排查、模型部署等等。已更新的文章根据篇章目录查看，相关代码实现开源在：https://github.com/wellinxu/nlp_store ，更多内容关注知乎专栏（或微信公众号）：NLP杂货铺。

基本依赖
自寻坑路
TensorRT

'UFF'模型转换错误

BERT in TensorRT

'time out'demo模型特难下载
顺利运行demo推断
'additional_dict'模型中包含训练时的参数
'cls_squad_output_weights'模型中参数名称不一致
'ascii'中文读取报错
模型结构与demo不同
'nan'batch大小的配置问题
tensorflow与tensorRT的推断时间对比

tensorRT server for BERT

'segmentation'cuda与flask的冲突
tensorrt模型文件的确认
'deserialize'tensorrtserver版本问题
'CustomEmbLayerNormPluginDynamic'插件缺失
config.pbtxt配置文件不对

Client

tensorrtserver.api下载安装
'dimension'batch大小的配置问题

上线

Serialization Error

总结
参考

基本依赖

[ ] python
[ ] git
[ ] gpu
[ ] nvidia-docker
[ ] bert的checkout模型

自寻坑路

那天接口调用方告诉我，我的接口超时了要优化下。从使用BERT开始我就知道，总会有这一天的，而现在终于来了。在一开始部署服务的时候，就直接上了GPU，对于一般调用，虽然慢了点，但也不至于超时。但被新业务调用后，处理的样本量明显增加，特别容易出现超时的现象。没办法，自己约的，呸，自己选的路，再艰难也得先看看别人趟的水，然后再决定走不走。既然需求已经提了，我当然立刻行动起来，根据项目特性，先优化了一波流程，让整体速度提升了4倍左右，之前timeout的样例都通过了，先部署起来给下游用。虽然提升了4倍速度，但大样本依然游走在timeout的边缘，如果遇到更大的样本，就...。
为了看看前人趟的水，就来知乎进行了搜索，然后得到了下面几篇文章。

《从零开始学习自然语言处理(NLP)》-BERT推理加速实践（6）【1】
《从零开始学习自然语言处理(NLP)》-BERT模型推理加速总结（5）【2】
加速 BERT 模型有多少种方法？从架构优化、模型压缩到模型蒸馏，最新进展详解！【3】
NVIDIA发布TensorRT 6，突破BERT-Large推理10毫秒大关【4】

从结果看，有这么几种方式：缩短max_seq，合并请求组成大batch，替换模型（蒸馏/缩减层数等），换成float16精度，使用tensorflow的xla,使用tensorRT。缩短max_seq首先被排除了，因为这个项目在处理过程中，要对所有的文本进行处理，缩短seq等于增加了样本量，所以不适合我们项目。组成大batch也不行，我们的一次请求都至少40个样本，在样本量超过40之后，batch的增加与时间的增加基本上是线性的，所以也不适合。至于替换模型，之前尝试过tiny版的albert,速度肯定有提升，但准确率降了5个点，接受不了。后面的几个方法，tensorRT同时包含了所有优点，so，基于tensorRT部署BERT服务，坑从此开始。

TensorRT

tensorRT【5】是什么，不知道，没听过，不管了，先按照说明【6】把tensorrt安装下，在我tensorflow14的docker容器中一顿操作，哎，木报错，顺利安装完。然后需要把模型转为tensorRT形式，顺利找到转tensorflow模型的文档【7】，当然checkout模型是不能直接转的，事先要转为‘frozen TensorFlow model’格式，这个在【7】中也有提示，也可以自行百度/google寻找教程。

'UFF'模型转换错误

顺利转完模型后，需要转为uff文件，这时坑开始来了，出现了一个它认识我我不认识它的错误（这个错误到写文章时都没解决，有一个原因是后来没有走这条路）。只能再去找前人，文章【8】中显示，nvidia专门对BERT进行过优化，是有demo的，一翻波折之后，终于找到了官方demo。

BERT in TensorRT

在tensortRT的官方github上一开始并没有找到BERT的demo，后来发现在5.1与6.0分支【9】中都有BERT的demo，聪明的我认为新版本应该更好些（也许选择5.1，会少走一些弯路，但我不想再去找坑了），于是重新开始了寻坑之旅，这样第一行代码出现了：

git clone -b release/6.0 https://github.com/NVIDIA/TensorRT.git

根据【9】中python的提示，依次运行：

cd TensorRT/demo/BERT
sh python/create_docker_container.sh

'time out'demo模型特难下载

在TensorRT/demo/BERT/python目录下有readme说明文档，根据说明准备环境，第一步克隆代码已经完成，第二步建立镜像，时间稍微有点久，但还算顺利，第三步在镜像中编译插件/下载调试过的模型（因为是完整的demo，所以模型都是有demo的），其中模型下载可以根据需求选择base/large，max_seq的大小，以及float32/float16精度,插件编译得还比较顺利，就是模型下载得太慢了，十几k每秒，也不知道多大，那就等着，自己再去看看不知道是什么的TensorRT。直到下班，还没下好，没事明天早上再来看，结果第二天过来看，告诉我下载失败，好吧重新下，5个小时后又下载失败，好吧重新下......终于下载了3天得到了base_fp16_384的模型。反应慢的我这时候想起，我们项目中的长度是128，最好弄个128的模型，这样方便对比推断的时间，所以在后台重新下载128的模型，自己先用384模型跑完demo的后续流程（事实上，一直到部署结束，都没能再成功下载过一个模型，想想当时能成功下载384的模型，真的是幸运）。

顺利运行demo推断

然后根据文档，顺利执行了"Building an Engine"与"Running Inference"两个步骤，运行时间有些波动，最快大概再2.4ms一个，鉴于长度是384，与官方说的128长度2.2ms一个基本吻合了。这个时候的我，仿佛已经完成了模型加速，幸福的表情难以言表。

'additional_dict'模型中包含训练时的参数

虽然说128长度的demo模型一直下不成功，但发现384的模型其实就是个checkout形式的tensorflow模型，所以就直接拿我们自己训练好的checkout模型来转换。开始转化模型：

python python/bert_builder.py -m /workspace/models/fine-tuned/bert_tf_v2_base_fp32_128_v2/model.ckpt-6001 -o bert_base_128.engine -b 1 -s 128 -c /workspace/models/fine-tuned/bert_tf_v2_base_fp16_128_v2

在短暂的等待之后，就迎来了ERROR!

报错表示，在模型数据加载过程中，出现了不支持的数据类型，于是在bert_builder.py的252/253行加入了代码：

print(pn)    # 打印参数名称
print(type(tensor))    # 打印参数类型

再次运行后发现，一个叫"signal_early_stopping/STOP"的参数是布尔形式，这是在训练过程中用到的，所以将237行改为：

param_names = [key for key in sorted(tensor_dict) if 'early_stop' not in key and 'adam' not in key and 'global_step' not in key and 'pooler' not in key]

'cls_squad_output_weights'模型中参数名称不一致

顺利解决了上面的bug，再次运行后，比上次多等待了一眨眼，就得到了一个全新的ERROR！

报错表示，参数中没有一个叫"cls_squad_output_weights"的，demo是squad的一个样例，而我的是二分类下游任务，最后输出层参数名称不一致，根据上面打印的参数名字（不同的下游任务或者训练代码，最后层的参数名字都可能不同），将217/218行代码修改为：

W_out = init_dict["output_weights"]
B_out = init_dict["output_bias"]

然后再次运行，没有报错，进入了相对较长的等待。当"Saving Engine to bert_base_128.engine Done."出现的时候，这个时候的我，仿佛已经完成了模型加速，幸福的表情难以言表。

'ascii'中文读取报错

赶紧以葫芦画瓢（将文档与问题都用文件的形式提供），也运行一次推断：

python python/bert_inference.py -e bert_base_128.engine -pf "p.txt" -qf "q.txt" -v /workspace/models/fine-tuned/bert_tf_v2_base_fp32_128_v2/vocab.txt

很快，就报了一个"ascii"编码的错误（因为读取中文的缘故），在百度之后，执行了下面一行得以解决(本容器中可使用C.UTF-8)：

export LANG=C.UTF-8

继续运行后，得到"Running inference in 437.362 Sentences/Sec"(2.286ms)，后面还有一个squad相关的错误，直接被我忽略了。这个时候的我，仿佛已经完成了模型加速，幸福的表情难以言表。

模型结构与demo不同

后面我修改了bert_inference.py文件（根据自己需要，自行修改），打印出模型分类结果，发现结果是[128,2,1,1]维度的，并且没有做最后的softmax操作，获取第一行数据并softmax，然后发现，结果是错的，什么鬼，摔！
这个错误让我茶饭不思（那两天吃得可好了），以为距离模型加速成功就一步之遥（其实还有好远），结果结果是错的，根本不能用！对着代码查这查那，用着google找这找那，丝毫没有头绪。一股神秘的力量，让我回去看了tensorflow训练BERT下游任务的代码，在我粗略的论文阅读中，以及网络/同事的介绍中，得到的信息都是，取BERT最后一层[CLS]的编码直接进行下游二分类训练(单层全连接+softmax)。但代码中却是，先经过一层全连接+tanh，再接全连接+softmax。根据代码得知，第一层的两参数名叫"bert_pooler_dense_kernel"，"bert_pooler_dense_bias"，所以将bert_build.py原237行修改为：

param_names = [key for key in sorted(tensor_dict) if 'early_stop' not in key and 'adam' not in key and 'global_step' not in key]

在tensorrt-api文档【10】的帮助下，将bert_build.py中的squad_output函数修改为：

def squad_output(prefix, config, init_dict, network, input_tensor):
     """
     Create the squad output
     """
      
     idims = input_tensor.shape
     assert len(idims) == 5
     B, S, hidden_size, _, _ = idims
      
     p_w = init_dict["bert_pooler_dense_kernel"]
     p_b = init_dict["bert_pooler_dense_bias"]
     #这里其实可以直接取[CLS]的向量进行后续运算，但是没能实现相关功能，就计算了所有的
     pool_output = network.add_fully_connected(input_tensor, hidden_size, p_w, p_b)
     pool_data = pool_output.get_output(0)
     tanh = network.add_activation(pool_data, trt.tensorrt.ActivationType.TANH)
     tanh_output = tanh.get_output(0)
      
     W_out = init_dict["output_weights"]
     B_out = init_dict["output_bias"]
      
     W = network.add_constant((1, hidden_size, 2), W_out)
     dense = network.add_fully_connected(tanh_output, 2, W_out, B_out)               
     set_layer_name(dense, prefix, "dense")
     return dense

在相对较长的等待中，获得了新的engine模型。在运行了自己修改的推断脚本后，得到了正确结果（使用了fp16,所以最后结果在万分位上有所不同，但基本一致了）。这个时候的我，仿佛已经完成了模型加速，幸福的表情难以言表。

'nan'batch大小的配置问题

为了对比下速度，多预测了几个样本，发现报错，可能是因为构建engine的时候batch设置的是1，重新设置为20，再次运行：

python python/bert_builder.py -m /workspace/models/fine-tuned/bert_tf_v2_base_fp32_128_v2/model.ckpt-6001 -o bert_base_128.engine -b 20 -s 128 -c /workspace/models/fine-tuned/bert_tf_v2_base_fp16_128_v2

得到新模型后，推断多个样本时，依然报错，得到的都是nan结果。什么情况，这个batch size参数是摆设吗!（后来我查看过5.1分支的代码，配置有所不同，也许在5.1分支上，直接使用batch size是有作用的）天知道在经过怎样的过程之后，发现engine构建过程中，有过配置设置：

bs1_profile = builder.create_optimization_profile()
set_profile_shape(bs1_profile, 1)
builder_config.add_optimization_profile(bs1_profile)

为了模型速度，构建过程中只设置了batch size为1，8以及参数值这三个。在google一翻之后，一位大佬说，再读取engine之后，设置使用第二个配置（就是以传入参数为batch size的配置）就可以了，也就是添加一行代码"context.active_optimization_profile = 1"，很有道理，一运行发现，"out of index"!还是要暴力处理，后来直接注释了其他两个配置，终于运行成功了，batch size在[1,20] 之间都得到了正确结果。这个时候的我，仿佛已经完成了模型加速，幸福的表情难以言表。

tensorflow与tensorRT的推断时间对比

然后就得到了一开始的那张时间对比表，tensorflow是1.14版，float32精度，使用estimator进行预测，tensorrt是6.0版，直接使用python调用：

max_seq	batch_size	1	5	10	20	40	100
128	Tensorflow(ms)	20	24	45	76	130	300
128	tensorRT(ms)	2.2	5.2	7.2	10.1	19	51

tensorRT server for BERT

'segmentation'cuda与flask的冲突

python调用成功后，就简单把bert_inference.py程序改成了一个python服务程序，顺利启动服务之后，调用服务接口，惊喜来了：

一个光秃秃的提示（报错位置都没有），反正没得结果。后来一查百度，说这个多为内存不当操作造成，这让我一个半路出家的程序员怎么办！后来仔细排查了报错位置，发现是cuda报出的错，结合这个信息，在面向google编程之后，得知pycuda上下文与http上下文有一些冲突（超过我知识范围了），在初始化cuda的时候不能使用autoinit,得调用一次init一次，

#import pycuda.autoinit #注释掉自动初始化
# 初始化cuda
cuda.init()
device = cuda.Device(0)
ctx = device.make_context()

# 中间所有处理程序

# 结束上下文
ctx.pop()

如上所示，每次调用的时候都得先初始化，经过这样的修改以后，果然跑通了。但耗时了3000ms一个样本，时间足足多了上百倍，还能不能好好地玩耍了。这还不是最关键的，最关键的是，所有的返回结果全部变成了[0.5,0.5],好吧在服务中用python直接调用模型是行不通了（并不是真正的行不通，只是我这个渣渣，cuda也不懂，解决不了这两个问题）。

tensorrt模型文件的确认

这个时候，只能去使用tensorrt server【11】了。根据【11】文档里所说，只要求一个配置文件和模型文件，就可以启动相应的docker服务了。但发现这个server能识别的tensorrt模型文件是一个以.plan结尾的文件，但我只有一个以.engine结尾的模型文件。后面我查阅了百度/gooogle相关文件，都没找到如何将engine转为plan文件，后来只找到一个在tensorrt6.0已经被弃用的方法"write_engine_to_file()",但这不能用，而且也不知道写进去的文件是啥类型。后来我搜索了【5】中所有的"plan"字符，终于找到一句话：

Write out the inference engine in a serialized format. This is also called a plan file.

这句话的意思应该就是说，engine跟plan是一个东东（后面确实读取成功了，应该是一个东西，当时还是有猜的成分）。

'deserialize'tensorrtserver版本问题

按照【11】中所讲，我认为下载官方提供的docker来部署最为方便，最主要的是，我发现服务器上已经有一个tensorrt server的镜像了，应该是同事之前下的，下镜像的过程都省掉了。根据【11】中Model相关内容，将模型文件重命名为model.plan，将配置文件修改为（有问题，后面要改）：

name: "bert"
platform: "tensorrt_plan"
max_batch_size: 20
input [
    {  
      name: "input0"
      data_type: TYPE_INT32
      dims: [128]
    }, 
    {  
      name: "input1"
      data_type: TYPE_INT32
      dims: [128]
    }, 
    {  
      name: "input2"
      data_type: TYPE_INT32
      dims: [128]
     }  
]      
output [
    {  
      name: "output0"
      data_type: TYPE_FP16
   dims:[128,2]                                                                                                         
    }  
]

根据【11】中的命令，直接运行：

NV_GPU=1 nvidia-docker run --rm --shm-size=1g --ulimit memlock=-1 --ulimit stack=67108864 -p50014:8000 -p50015:8001 -p50016:8002 -v /自己的路径/models:/models nvcr.io/nvidia/tensorrtserver:19.08-py3 trtserver --model-repository=/models

直接一个报错"trtserver: unrecognized option '--model-repository=/models'"，根据提示，将参数名"--model-repository"改为了"--model-store"，然后再运行，就得到了又一个错误：

在一个不知道在哪里的文件，报了一个c++的错误，一筹莫展。根据提示，好像是batch size的问题，也可能是engine文件并不是plan文件导致读错了，也可能是版本问题（tensorrt server跟tensorrt的版本不统一【12】）。针对前面两种情况，试了各种姿势，依然是报这个错，没办法，只能重新下载个docker镜像了。

docker pull nvcr.io/nvidia/tensorrtserver:19.09-py3

当然下载不会一帆风顺的，会有权限错误提示，在【11】中也多次提到，下载docker容器得先有NGC 权限【13】。

'CustomEmbLayerNormPluginDynamic'插件缺失

两个小时之后，下载完毕，然后运行：

NV_GPU=1 nvidia-docker run --rm --shm-size=1g --ulimit memlock=-1 --ulimit stack=67108864 -p50014:8000 -p50015:8001 -p50016:8002 -v /自己的路径/models:/models nvcr.io/nvidia/tensorrtserver:19.09-py3 trtserver --model-repository=/models

大概等了2秒钟，有点开心，已经过了上个bug出现的时间，又过了1秒，果然来了个新bug：

读取"CustomEmbLayerNormPluginDynamic"插件错误，这个插件有点眼熟，在bert_build.py文件里面出现过，在咨询了一波google之后，发现一个类似的问题【14】，在根据bert_build.py文件，将libbert_plugins.so/libcommon.so(在【9】中build文件夹中，编译之后的)文件拷进docker容器里。进入容器：

NV_GPU=1 nvidia-docker run -it --shm-size=1g --ulimit memlock=-1 --ulimit stack=67108864 -p50014:8000 -p50015:8001 -p50016:8002 -v /自己的路径/models:/models nvcr.io/nvidia/tensorrtserver:19.09-py3 /bin/bash

然后运行：

export LD_PRELOAD=/opt/tensorrtserver/libbert_plugins.so:/opt/tensorrtserver/libcommon.so

so文件位置是自己确定的，然后运行：

trtserver --model-repository=/models

config.pbtxt配置文件不对

这次运行时间更长了，好开心，然后：

根据这个提示，是配置文件有问题，修改后又报错再修改再报错，这样循环几次后，配置文件变成了：

name: "bert"
platform: "tensorrt_plan"
max_batch_size: 20
input [
    {
      name: "input_ids"
      data_type: TYPE_INT32
      dims: [128]
    },
    {
      name: "segment_ids"
      data_type: TYPE_INT32
      dims: [128]
    },
    {
      name: "input_mask"
      data_type: TYPE_INT32
      dims: [128]
    }
]
output [
    {
      name: "cls_dense"
      data_type: TYPE_FP32
      dims: [128, 2, 1, 1]
    }
]

再次运行之后，一直没报错，好像成功了，根据【11】中状态检查，运行了下：

curl localhost:50014/api/status

然后得到了：

哦耶，好像加载成功了，这个时候的我，仿佛已经完成了模型加速，幸福的表情难以言表。

Client

tensorrtserver.api下载安装

看似tensorrt的服务已经成功部署了，现在就需要在我自己的服务内部调用成功了。根据【11】当中关于client的介绍，client的环境可以自己打镜像/编译/下载镜像/下载编译好的结果，琢磨着下载编译好的结果最简单了，所以在【15】中找到了自己需要的版本，直接运行：

wget https://github.com/NVIDIA/tensorrt-inference-server/releases/download/v1.6.0/v1.6.0_ubuntu1804.clients.tar.gz
tar -zxvf v1.6.0_ubuntu1804.clients.tar.gz

然后安装一下：

cd python
pip install tensorrtserver-1.6.0-py2.py3-none-linux_x86_64.whl

'dimension'batch大小的配置问题

有了依赖环境之后，根据【11】中的python api，成功地将模型调用添加到我的服务中去，最后调用测试，这个时候的我，仿佛已经完成了模型加速，幸福的表情难以言：

根据报错提示，是batch size不对，将样本数量换成20后，果然运行正确了。为什么直接用python调用模型的时候，样本数只要小于batch size就可以了！后来我将bert_build.py文件的"set_profile_shape"函数改为了：

def set_profile_shape(profile, batch_size):
    maxshape = (batch_size, S)
    minshape = (1, S)
    optshape = (batch_size, S) # 这个batch的大小在最大和最小之间就可以，可以相等
    profile.set_shape("input_ids", min=shape, opt=shape, max=shape)
    profile.set_shape("segment_ids", min=shape, opt=shape, max=shape)
    profile.set_shape("input_mask", min=shape, opt=shape, max=shape)

其中"set_shape"等函数的含义可以在文档【10】中找到，修改后，又将全流程走了一遍，构建engine，部署tensorrt server，运行client，然后我真的完成了了模型加速，幸福的表情一言难尽！最后测试20个样本调用服务的时间是15ms左右，比python直接调用延迟了5ms左右。

上线

Serialization Error

果然，我还是笑得太早，太年轻了。刚把服务推到测试上，就来了惊喜：

一看错误，模型加载错误，什么情况，我不是已经跑通了吗？经过一系列查询之后，发现tensorrt在不同GPU（主要根据计算能力分类）上编译的模型是不能通用的，我之前在调研机上跑的，GPU型号是V100（计算能力7.0），而测试机上是P4（计算能力6.1），在V100上编译的模型不能再P4使用，根据github上前辈提示，在CMakeLists.txt文件的第21行添加：

-gencode arch=compute_60,code=sm_60 \

这样就可以在P4上进行编译了，后来验证P4上的模型可以在P100的机器上跑通。当然编译的过程显然不能一帆风顺，中间出现了out of memory的问题，后来测试发现，大概需要5800M的显存才能编译成功。线上的GPU是P100的，顺利运行了新编译的模型，心里终于微微一笑；bert在V100上的运行速度大概是P100的5倍，而且nvidia主要在T4跟V100上优化bert推断的，而且模型转换时P100上不能使用float16精度模式（只能用float32），所以新编译的模型效率提高有限；想到新编译的模型比tensorflow的效率只提高了30%左右，感觉之前的努力突然不香了。

总结

压力/需求使人进步，如果没有"time out"报错，我在部署完gpu版的BERT模型就结束了，这不，硬着头皮也要上，从完全没有听过tensorrt到成功部署bert，大概用了两周多时间。虽然现在对tensorrt依然只是了解皮毛，对cuda编程更是两眼一抹黑，但没关系，有了这个开端，后面慢慢学。
在部署过程中，百度/google/知乎各种找教程，都没有找到部署BERT的详细过程，也为了提升自身，写下自己遇到的坑，以及部分解决办法，供大家一起交流进步。

参考（阅读原文可查看链接）

【1】《从零开始学习自然语言处理(NLP)》-BERT推理加速实践（6）
【2】《从零开始学习自然语言处理(NLP)》-BERT模型推理加速总结（5）
【3】加速 BERT 模型有多少种方法？从架构优化、模型压缩到模型蒸馏，最新进展详解！
【4】NVIDIA发布TensorRT 6，突破BERT-Large推理10毫秒大关
【5】What Is TensorRT?
【6】tensorRT的debian方式安装
【7】使用python将tensorflow模型转为tensorRT模型
【8】Real-Time Natural Language Understanding with BERT Using TensorRT
【9】tensorRT6.0分支BERT官方demo
【10】tensorrt-api
【11】NVIDIA TensorRT Inference Server
【12】TensorRT inference server documentation versions
【13】NGC guide
【14】CustomEmbLayerNormPluginDynamic插件加载问题
【15】tensorrt-inference-server/releases
【16】DeepLearningExamples