书接上回
在之前的文章《GPU服务器初体验:从零搭建Pytorch GPU开发环境 》中,我通过Github上一个给新闻标题做分类的Bert项目,演示了Pytorch模型训练与预测的过程。我其实也不是机器学习的专业人士,对于模型的结构、训练细节所知有限。作为后台开发而非算法工程师,我更关注的是模型部署的过程。
小议Online Serving
前文中,我们虽然有了python版预测脚本,但在实际生产过程中,我们还需将模型给服务化,能对外通过接口提供在线预测、推理的能力。这一过程称为Online Serving 或者 Online Inference,即在线Serving、在线推理。俗称模型部署。
若将pyhton代码服务化,在性能方面是不能满足要求的,无法做到低延时和高吞吐。因此生产环境一般使用编译型语言来加载模型提供预测推理服务。这个领域最常用的编程语言就是C++,比如TensorFlow配套的TF-Serving。但Pytorch官方没有提供线上Serving的方案,常见的解决方案是将Pytorch模型转为ONNX模型,再通过ONNX模型的服务化方案来部署到线上。
ONNX 与 ONNX Runtime
ONNX是Open Neural Network Exchange的缩写,打开它的官网:onnx.ai/ 几个大字赫然而出:
Open Neural Network Exchange The open standard for machine learning interoperability
开放神经网络交换(格式),机器学习互操作性的开放标准。
ONNX是2017年9月由微软与Facebook、AWS合作推出的开放的神经网络交换格式。致力于将不同模型转换成统一的ONNX格式,然后再通过统一的方案完成模型部署。不由得让人想起那句计算机领域的经典论断:
在软件工程中,没有一个中间层解决不了的问题
没错,ONNX也是一种“中间层”的概念,好比LLVM的IR,将编译器的工程分层,一层开放给不同编程语言实现,一层对接不同的硬件OS,中间通过IR串联。ONNX作为中间层,的一头对接不同的机器学习模型框架,另外一头对接的是不同的编程语言(C++、Java、C#、Python……)、不同OS(windows、Linux……)、不同计算引擎(CPU、CUDA……)的模型部署方案。各种机器学习框架产出的模型,只需要转换成ONNX格式,就自然获得了在多种平台上使用多种编程语言做在线推理的能力。
准确的说ONNX的部署和服务化是由另外一个项目完成的,即ONNX Runtime。它有独立的产品品牌和官网:onnxruntime.ai/ 。当然它也是微软主导的项目。Onnx Runtime其实不只是单纯地完成模型的部署,也会对模型推理过程有一些优化。
回顾Pytorch预测脚本
先回顾一下前文中的Pytorch模型预测脚本pred.py,代码是从这个issue直接拿来主义的:单条文本数据的预测代码 #72 感谢这位网友。 另外因为我们的模型是GPU训练的,所以对代码做了修改,他的代码是是CPU模型做预测的。
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资料: pytorch.org/docs/stable… 针对我们的得模型,我们可以这样写出大致的导出脚本 to_onnx.py:
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 | 动态维度设置,不设置即只支持固定维度的参数。本例子其实可以不设置,因为我们传入的参数都是自己调整好维度的。 |
args参数的探讨
args用于标识模型输入参数的shape。这个可以好好谈谈一下。
参数错误?
回顾一下前面的pytorch模型预测脚本,build_predict_text()函数会对一段文本处理成模型的三个输入参数,所以它返回的对象肯定是符合模型输入shape的。:
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, mask
torch.onnx.export()调用的时候,其实只关心形状,而不关心内容。所以我们可以直接改成这样:
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()
... ...
但如果你这样调用了,会报错:
File "https://b2.7b2.com/home/guodong/github/guodong/Bert-Chinese-Text-Classification-Pytorch/to_onnx.py", line 64, in <module>
torch.onnx.export(model,
File "https://b2.7b2.com/home/guodong/miniconda3/envs/onnx_gpu/lib/python3.10/site-packages/torch/onnx/utils.py", line 504, in export
_export(
File "https://b2.7b2.com/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 "https://b2.7b2.com/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 "https://b2.7b2.com/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 "https://b2.7b2.com/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 "https://b2.7b2.com/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 "https://b2.7b2.com/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 "https://b2.7b2.com/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 "https://b2.7b2.com/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 "https://b2.7b2.com/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 "https://b2.7b2.com/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中)
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存储的就是参数输入参数。 其实很多人都遇到过这个问题,比如:
应该是tensor.onnx.export()内部把args这个tuple给unpack(展开)了,所以函数参数变多了。
解决方法如下:
写法一
可以给它再套上一层。
if __name__ == '__main__':
args = build_arg1()
torch.onnx.export(model,
(args,),
'model.onnx',
... ...
如果你实在不想给args再套一层,可以让build_args1返回list,实测也能解决问题。
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]
进而也可以转换成另外一种写法:
写法二
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()则是整型,且参数和randn()不一样。 randint()的完整声明如下:
torch.randint(low=0, high, size, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False)
忽略带默认值的参数,只需关注:
torch.randint(low=0, high, size)
low和high就是随机整数的范围,size是这个Tensor的shape,需要用元组表示。比如(1, pad_size) 表示的行数为1,列数为pad_size。 值得一提的是,seq_len的shape不是二维的,它是标量,只有一维。但如果你写成
seq_len = torch.randint(1, 10, (1)).cuda()
则报错,需要写成
seq_len = torch.randint(1, 10, (1,)).cuda()
这倒不是pytorch或onnx的坑,而是python语言的坑,因为当元组只有一个元素的时候,它其实会直接退化成这个元素的类型。 打开一个python交互式命令行一试便知:
>>> 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安装命令如下:
pip install onnxruntime-gpu -i https://pypi.tuna.tsinghua.edu.cn/simple
如果去掉-gpu,则安装的也是CPU版本。
执行一下下面脚本,检查一下是否有安装成功:
import onnxruntime as ort
print(ort.__version__)
print(ort.get_device())
在我的环境上会输出:
1.13.1
GPU
创建InferenceSession对象
onnxruntime的python API手册在这:onnxruntime.ai/docs/api/py…
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()方法,它的参数声明如下:
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变量表示参数的名称。 写个小代码测试一下:
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数组。可以在网上找的转换的代码:
def to_numpy(tensor):
return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
然后就可以:
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拆成两部分:
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预测脚本可以这么写:
#!/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))
最终输出:
天问一号着陆火星一周年 is science
耗时对比
接下来,我们可以对比一下onnxruntime使用CUDA和之前直接pytorch CUDA做预测的耗时变化。 先写一些辅助函数:
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函数如下:
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')
最终结果:
杭州购房政策大松绑 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模型完成一个在线预测服务。但限于文章篇幅,我们下次再聊!请大家继续关注。