在软件工程中，没有一个中间层解决不了的问题序言
小议Online Serving
在之前的文章 《GPU服务器初体验：从零搭建Pytorch GPU开发环境》 中，我通过Github上一个给新闻标题做分类的Bert项目，演示了Pytorch模型训练与预测的过程。我其实也不是机器学习的专业人士，对于模型的结构、训练细节所知有限，但作为后台开发而非算法工程师，我更关注的是模型部署的过程。
在前文中，我们虽然有了python版预测脚本，但在实际生产过程中，我们还需将模型给服务化，能对外通过接口提供在线预测、推理的能力。这一过程称为Online Serving 或者 Online Inference，即在线Serving、在线推理。俗称模型部署。
若将pyhton代码服务化，在性能方面其实是不能满足要求的，无法做到低延时和高吞吐。因此生产环境一般使用编译型语言来加载模型提供预测推理服务。这个领域最常用的编程语言就是C ，比如TensorFlow配套的TF-Serving。但Pytorch官方没有提供线上Serving的方案，常见的解决方案是将Pytorch模型转为ONNX模型，再通过ONNX模型的服务化方案来部署到线上。
ONNX 与 ONNX Runtime
打开onnx的官网：https://onnx.ai/
几个大字赫然而出：
Open Neural Network Exchange The open standard for machine learning interoperability
开放神经网络交换（格式），机器学习互操作性的开放标准。
ONNX是2017年9月由微软与Facebook、AWS合作推出的开放的神经网络交换格式。致力于将不同模型转换成统一的ONNX格式，然后再通过统一的方案完成模型部署。
没错，这也是一种“中间层”的概念，好比LLVM的IR，将编译器的工程分层，一层开放给不同编程语言实现，一层对接不同的硬件OS，中间通过IR串联。ONNX作为中间层，的一头对接不同的机器学习模型框架，另外一头对接的是不同的编程语言（C 、Java、C#、Python……）、不同OS（windows、Linux……）、不同算力设备（CPU、CUDA……）的模型部署方案。
准确的说ONNX的部署和服务化是由另外一个项目完成的，即ONNX Runtime。它有独立的产品品牌和官网：https://onnxruntime.ai/ 。当然它也是微软主导的项目。Onnx Runtime其实不只是单纯地完成模型的部署，也会对模型推理过程有一些优化。
回顾Pytorch预测脚本
先回顾一下前文中的Pytorch模型预测脚本pred.py，代码是从这个issue直接拿来主义的：单条文本数据的预测代码 #72 感谢这位网友。
另外因为我们的模型是GPU训练的，所以对代码做了修改，他的代码是是CPU模型做预测的。
代码语言：python代码运行次数：0复制import torch
from importlib import import_module
key = {
0: 'finance',
1: 'realty',
2: 'stocks',
3: 'education',
4: 'science',
5: 'society',
6: 'politics',
7: 'sports',
8: 'game',
9: 'entertainment'
}
model_name = 'bert'
x = import_module('models.' model_name)
config = x.Config('THUCNews')
model = x.Model(config).to(config.device)
model.load_state_dict(torch.load(config.save_path, map_location='cpu'))
def build_predict_text(text):
token = config.tokenizer.tokenize(text)
token = ['[CLS]'] token
seq_len = len(token)
mask = []
token_ids = config.tokenizer.convert_tokens_to_ids(token)
pad_size = config.pad_size
if pad_size:
if len(token) < pad_size:
mask = [1] * len(token_ids) ([0] * (pad_size - len(token)))
token_ids = ([0] * (pad_size - len(token)))
else:
mask = [1] * pad_size
token_ids = token_ids[:pad_size]
seq_len = pad_size
ids = torch.LongTensor([token_ids]).cuda() # 改了这里，加上.cuda()
seq_len = torch.LongTensor([seq_len]).cuda() # 改了这里，加上.cuda()
mask = torch.LongTensor([mask]).cuda() # 改了这里，加上.cuda()
return ids, seq_len, mask
def predict(text):
"""
单个文本预测
"""
data = build_predict_text(text)
with torch.no_grad():
outputs = model(data)
num = torch.argmax(outputs)
return key[int(num)]
if __name__ == '__main__':
print(predict("备考2012高考作文必读美文50篇(一)"))把Pytorch模型导出成ONNX模型
torch.onnx.export()基本介绍
pytorch自带函数torch.onnx.export()可以把pytorch模型导出成onnx模型。官网API资料：
https://pytorch.org/docs/stable/onnx.html#torch.onnx.export
针对我们的得模型，我们可以这样写出大致的导出脚本 to_onnx.py：
代码语言：python代码运行次数：0复制import torch
from importlib import import_module
model_name = 'bert'
x = import_module('models.' model_name)
config = x.Config('THUCNews')
model = x.Model(config).to(config.device)
model.load_state_dict(torch.load(config.save_path))
def build_args():
pass #... 先忽略
if __name__ == '__main__':
args = build_arg()
torch.onnx.export(model,
args,
'model.onnx',
export_params = True,
opset_version=11,
input_names = ['ids','seq_len', 'mask'], # the model's input names
output_names = ['output'], # the model's output names
dynamic_axes={'ids' : {0 : 'batch_size'}, # variable lenght axes
'seq_len' : {0 : 'batch_size'},
'mask' : {0 : 'batch_size'},
'output' : {0 : 'batch_size'}})对export函数的参数进行一下解读：
参数 | 解读 |
|---|---|
model | 加载的pytorch模型的变量 |
args | 指的是模型输入的shape（形状） |
'model.onnx' | 导出的onnx模型的文件名 |
export_params | 是否导出参数 |
opset_version | ONNX的op版本，这里用的是11 |
input_names | 模型输入的参数名 |
output_names | 模型输出的参数名 |
dynamic_axes | 动态维度设置，不设置即只支持固定维度的参数。本例子其实可以不设置，因为我们传入的参数都是自己调整好维度的。 |
其实对于Bert模型的输入而言，seq_len是不需要的。可以阅读一下模型的forward()函数（前向传播的函数）定义，其实里面丢弃了seq_len。
args参数的探讨
args用于标识模型输入参数的shape。这个可以好好谈谈一下。
参数错误？
回顾一下前面的pytorch模型预测脚本，build_predict_text()函数会对一段文本处理成模型的三个输入参数，所以它返回的对象肯定是符合模型输入shape的。：
代码语言：python代码运行次数：0复制def build_predict_text(text):
token = config.tokenizer.tokenize(text)
token = ['[CLS]'] token
seq_len = len(token)
mask = []
token_ids = config.tokenizer.convert_tokens_to_ids(token)
pad_size = config.pad_size
if pad_size:
if len(token) < pad_size:
mask = [1] * len(token_ids) ([0] * (pad_size - len(token)))
token_ids = ([0] * (pad_size - len(token)))
else:
mask = [1] * pad_size
token_ids = token_ids[:pad_size]
seq_len = pad_size
ids = torch.LongTensor([token_ids]).cuda()
seq_len = torch.LongTensor([seq_len]).cuda()
mask = torch.LongTensor([mask]).cuda()
return ids, seq_len, masktorch.onnx.export()调用的时候，其实只关心形状，而不关心内容。所以我们可以直接改成这样：
代码语言：python代码运行次数：0复制def build_args1():
pad_size = config.pad_size
ids = torch.LongTensor([[0]*pad_size]).cuda()
seq_len = torch.LongTensor([0]).cuda()
mask = torch.LongTensor([[0]*pad_size]).cuda()
return [ids, seq_len, mask]
if __name__ == '__main__':
args = build_args1()
... ...但如果你这样调用了，会报错：
代码语言：txt复制File "/home/guodong/github/guodong/Bert-Chinese-Text-Classification-Pytorch/to_onnx.py", line 64, in <module>
torch.onnx.export(model,
File "/home/guodong/miniconda3/envs/onnx_gpu/lib/python3.10/site-packages/torch/onnx/utils.py", line 504, in export
_export(
File "/home/guodong/miniconda3/envs/onnx_gpu/lib/python3.10/site-packages/torch/onnx/utils.py", line 1529, in _export
graph, params_dict, torch_out = _model_to_graph(
File "/home/guodong/miniconda3/envs/onnx_gpu/lib/python3.10/site-packages/torch/onnx/utils.py", line 1111, in _model_to_graph
graph, params, torch_out, module = _create_jit_graph(model, args)
File "/home/guodong/miniconda3/envs/onnx_gpu/lib/python3.10/site-packages/torch/onnx/utils.py", line 987, in _create_jit_graph
graph, torch_out = _trace_and_get_graph_from_model(model, args)
File "/home/guodong/miniconda3/envs/onnx_gpu/lib/python3.10/site-packages/torch/onnx/utils.py", line 891, in _trace_and_get_graph_from_model
trace_graph, torch_out, inputs_states = torch.jit._get_trace_graph(
File "/home/guodong/miniconda3/envs/onnx_gpu/lib/python3.10/site-packages/torch/jit/_trace.py", line 1184, in _get_trace_graph
outs = ONNXTracedModule(f, strict, _force_outplace, return_inputs, _return_inputs_states)(*args, **kwargs)
File "/home/guodong/miniconda3/envs/onnx_gpu/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
return forward_call(*input, **kwargs)
File "/home/guodong/miniconda3/envs/onnx_gpu/lib/python3.10/site-packages/torch/jit/_trace.py", line 127, in forward
graph, out = torch._C._create_graph_by_tracing(
File "/home/guodong/miniconda3/envs/onnx_gpu/lib/python3.10/site-packages/torch/jit/_trace.py", line 118, in wrapper
outs.append(self.inner(*trace_inputs))
File "/home/guodong/miniconda3/envs/onnx_gpu/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
return forward_call(*input, **kwargs)
File "/home/guodong/miniconda3/envs/onnx_gpu/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1178, in _slow_forward
result = self.forward(*input, **kwargs)
TypeError: Model.forward() takes 2 positional arguments but 4 were given报错显示，forward()函数预期传入两个参数，但是实际传入了4个。
看一下模型的forward函数的定义（models/bert.py中)
代码语言：python代码运行次数：0复制class Model(nn.Module):
... ...
def forward(self, x):
context = x[0] # 输入的句子
mask = x[2] # 对padding部分进行mask，和句子一个size，padding部分用0表示，如：[1, 1, 1, 1, 0, 0]
_, pooled = self.bert(context, attention_mask=mask, output_all_encoded_layers=False)
out = self.fc(pooled)
return out函数定义确实是两个参数，一个self，一个x，x存储的就是参数输入参数。
其实很多人都遇到过这个问题:
https://github.com/pytorch/pytorch/issues/11456
https://github.com/onnx/onnx/issues/2711
应该是tensor.onnx.export内部把args这个tuple给unpack（展开）了，所以函数参数变多了。解决方法可以给它在套上一层。比如：
写法一
代码语言：python代码运行次数：0复制if __name__ == '__main__':
args = build_arg1()
torch.onnx.export(model,
(args,),
'model.onnx',
... ...写法二
如果你实在不想给args再套一层，可以让build_args1返回list，实测也能解决问题。
代码语言：python代码运行次数：0复制def build_args1():
pad_size = config.pad_size
ids = torch.LongTensor([[0]*pad_size]).cuda()
seq_len = torch.LongTensor([0]).cuda()
mask = torch.LongTensor([[0]*pad_size]).cuda()
return [ids, seq_len, mask]进而也可以转换成另外一种写法：
写法二
代码语言：python代码运行次数：0复制def build_args2():
pad_size = config.pad_size
ids = torch.randint(1, 10, (1, pad_size)).cuda()
seq_len = torch.randint(1, 10, (1,)).cuda() # 第三个参数中逗号不能少
mask = torch.randint(1, 10, (1, pad_size)).cuda()
return [ids, seq_len, mask]
if __name__ == '__main__':
args = build_args2()这里是用随机数来初始化torch.Tensor()，用randint，而没有用randn，是因为实测发现，不仅shape要对齐，数据类型也需要匹配。randn构造的Tensor是浮点型的。randint则是整型，也因为是整型，所以randint的参数和randn不一样。
randint的完整声明如下：
代码语言：python代码运行次数：0复制torch.randint(low=0, high, size, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False)忽略带默认值的参数，只需关注：
代码语言：python代码运行次数：0复制torch.randint(low=0, high, size)low和high就是随机整数的范围，size是这个Tensor的shape，需要用元组表示。比如(1, pad_size) 表示的行数为1，列数为pad_size。
值得一提的是，seq_len的shape不是二维的，它是标量，只有一维。但如果你写成
代码语言：python代码运行次数：0复制seq_len = torch.randint(1, 10, (1)).cuda()则报错，需要写成
代码语言：python代码运行次数：0复制seq_len = torch.randint(1, 10, (1,)).cuda()这倒不是pytorch或onnx的坑，而是python语言的坑，因为当元组只有一个元素的时候，它其实会直接退化成这个元素的类型。
打开一个python交互式命令行一试便知：
代码语言：txt复制>>> a=(1,2)
>>> type(a)
<class 'tuple'>
>>> a=(1)
>>> type(a)
<class 'int'>
>>> type(a)
<class 'tuple'>用ONNX Runtime做预测
好了，经过前面的步骤，顺利的话，已经得到一个onnx的模型文件model.onnx了，现在我们可以加载这个模型并执行预测任务。但我们不能一口气吃成一个胖子，在真正使用C 将ONNX模型服务化之前，我们还是需要先使用Python完成ONNX模型的预测，一方面是验证我们转换出来的ONNX确实可用，另一方面对后续我们换其他语言来服务化也有参考意义！
这个过程我们就需要用到ONNX Runtime的库：onnxruntime了。onnxruntime通常简称ort。
安装onnxruntime-gpu
onnxruntime不会随onnx一起安装，需要单独安装。因为我们整个实际都是基于GPU展开的，这里推荐用pip安装，因为conda似乎没有gpu的包，conda默认安装的是CPU版本。pip安装命令如下：
代码语言：shell复制pip install onnxruntime-gpu -i https://pypi.tuna.tsinghua.edu.cn/simple 如果去掉-gpu，则安装的也是CPU版本。
执行一下下面脚本，检查一下是否有安装成功：
代码语言：python代码运行次数：0复制import onnxruntime as ort
print(ort.__version__)
print(ort.get_device())在我的环境上会输出：
代码语言：shell复制1.13.1
GPU创建InferenceSession对象
onnxruntime的python API手册在这：https://onnxruntime.ai/docs/api/python/api_summary.html
onnxruntime中执行预测的主体是通过InferenceSession类型的对象进行的。InferenceSession 常用的构造参数只有2个，示例：
sess = ort.InferenceSession('model.onnx', providers=['CUDAExecutionProvider'])第一个参数就是模型文件的路径，第二个参数指定provider，它的取值可以是：
- CUDAExecutionProvider
- CPUExecutionProvider
- TensorrtExecutionProvider
顾名思义，CUDAExecutionProvider就是用GPU的CUDA执行，CPUExecutionProvider就是用CPU执行，TensorrtExecutionProvider是用TensorRT执行。没有安装TensorRT环境的话，即使指定它也不会生效，会退化成·CPUExecutionProvider·。
预测
预测过程就是InferenceSession对象调用run()方法，它的参数声明如下：
代码语言：python代码运行次数：0复制run(output_names, input_feed, run_options=None)- output_names – 输出的名字，可以为None
- input_feed – 字典类型
{ 输入参数名: 输入参数的值 } - run_options – 有默认值，可以忽略
它的返回值是一个list，list里面的值可以理解成是这段预测文本与每种分类的概率。我们再找到概率最大的分类就是最终结果了。
好了，现在唯一的问题就是构造第二个参数了。这是一个字典数据结构，key是参数的名称。我们可以通过InferenceSession的get_inputs()函数来获取。get_inputs()返回一个list，list中NodeArg类型的对象，这个对象有一个name变量表示参数的名称。
写个小代码测试一下：
代码语言：python代码运行次数：0复制a = [x.name for x in sess.get_inputs()]
print(a)['ids', 'mask']可以看到两个输入参数的名称ids和mask，其实就是我们导出ONNX模型的时候指定的输入参数名，前面我提到过seq_len其实没参与训练，所以不进模型。
好了，key搞定了，我们再来搞定value。
pytorch的预测过程中，我们通过 build_predict_text()把一段文本转换成了三个torch.Tensor。onnx模型的输入肯定不是torch中的Tensor。它只需要numpy数组即可。
偷懒的做法是，我们直接引入build_predict_text()，然后把Tensor类型转换成numpy数组。可以在网上找的转换的代码：
代码语言：python代码运行次数：0复制def to_numpy(tensor):
return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()然后就可以：
代码语言：python代码运行次数：0复制def predict(sess, text):
ids, seq_len, mask = build_predict_text(t)
print(type(ids))
input = {
sess.get_inputs()[0].name: to_numpy(ids),
sess.get_inputs()[1].name: to_numpy(mask),
}
outs = sess.run(None, input)
num = np.argmax(outs)
return key[num]其实这有点绕弯子了，我们其实不需要通过Tensor转numpy，因为Tensor是通过list转出来的，我们直接用list转numpy就可以了。我们先改一下pred.py将原先的build_predict_text拆成两部分：
代码语言：python代码运行次数：0复制def build_predict_text_raw(text):
token = config.tokenizer.tokenize(text)
token = ['[CLS]'] token
seq_len = len(token)
mask = []
token_ids = config.tokenizer.convert_tokens_to_ids(token)
pad_size = config.pad_size
# 下面进行padding，用0补足位数
if pad_size:
if len(token) < pad_size:
mask = [1] * len(token_ids) ([0] * (pad_size - len(token)))
token_ids = ([0] * (pad_size - len(token)))
else:
mask = [1] * pad_size
token_ids = token_ids[:pad_size]
seq_len = pad_size
return token_ids, seq_len, mask
def build_predict_text(text):
token_ids, seq_len, mask = build_predict_text_raw(text)
ids = torch.LongTensor([token_ids]).cuda()
seq_len = torch.LongTensor([seq_len]).cuda()
mask = torch.LongTensor([mask]).cuda()接着我们onnx的预测脚本（onnx_pred.py）中就可以直接调用build_predict_text_raw()，再把它的结果转numpy数组就可以了。注意，我们训练得到的Bert模型需要的是一个二维结构，所以和Tensor的构造方式一样，还需要再套上一层[] 。
好了，完整的onnx预测脚本可以这么写：
代码语言：python代码运行次数：0复制#!/usr/bin/env python
# coding=utf-8
import numpy as np
import onnxruntime as ort
import pred
def predict(sess, text):
ids, seq_len, mask = pred.build_predict_text_raw(text)
input = {
'ids': np.array([ids]),
'mask': np.array([mask]),
}
outs = sess.run(None, input)
num = np.argmax(outs)
return pred.key[num]
if __name__ == '__main__':
sess = ort.InferenceSession('./model.onnx', providers=['CUDAExecutionProvider'])
t = '天问一号着陆火星一周年'
res = predict(sess, t)
print('%s is %s' % (t, res))最终输出：
代码语言：txt复制天问一号着陆火星一周年 is science性能对比
模型转换为ONNX模型后，性能表现如何呢？接下来可以验证一下。 先写一些辅助函数：
代码语言：python代码运行次数：0复制def load_title(fname):
"""
从一个文件里加载新闻标题
"""
ts = []
with open(fname) as f:
for line in f.readlines():
ts.append(line.strip())
return ts
def batch_predict(ts, predict_fun, name):
"""
使用不同的预测函数，批量预测，并统计耗时
"""
print('')
a = time.time()
for t in ts:
res = predict_fun(t)
print('%s is %s' % (t, res))
b = time.time()
print('%s cost: %.4f' % (name, (b - a)))news_title.txt中有多条新闻标题，我们来让pytorch模型和onnx模型分别做一下预测，然后看耗时就可以了。
main函数如下：
代码语言：python代码运行次数：0复制if __name__ == '__main__':
model_path = './model.onnx'
cuda_ses = ort.InferenceSession('./model.onnx', providers=['CUDAExecutionProvider'])
ts = load_title('./news_title.txt')
batch_predict(ts, lambda t: predict(cuda_ses, t), 'ONNX_CUDA')
batch_predict(ts, lambda t: pred.predict(t), 'Pytorch_CUDA')最终结果：
代码语言：txt复制杭州购房政策大松绑 is realty
兰州野生动物园观光车侧翻事故新进展：2人经抢救无效死亡 is society
4个小学生离家出走30公里想去广州塔 is society
朱一龙戏路打通电影电视剧 is entertainment
天问一号着陆火星一周年 is science
网友拍到天舟五号穿云而出 is science
ONNX_CUDA cost: 0.0406
杭州购房政策大松绑 is realty
兰州野生动物园观光车侧翻事故新进展：2人经抢救无效死亡 is society
4个小学生离家出走30公里想去广州塔 is society
朱一龙戏路打通电影电视剧 is entertainment
天问一号着陆火星一周年 is science
网友拍到天舟五号穿云而出 is science
Pytorch_CUDA cost: 0.0888可以看到ONNX为Pytorch的耗时快了一倍。当然TensorRT的耗时应该会更低，不过这是后话了，本文暂且不表。
待续
好了，到此为止我们已经验证了转换后的ONNX模型可用性以及性能。下一步我们将使用C 来部署ONNX模型完成一个在线预测服务。但限于文章篇幅，我们下次再聊！请大家继续关注。
