实践演练Pytorch Bert模型转ONNX模型及预测

果冻虾仁 • 2023-01-02 • 云技术社区 • 278 阅读

在软件工程中，没有一个中间层解决不了的问题序言

小议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模型做预测的。

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

针对我们的得模型，我们可以这样写出大致的导出脚本 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	动态维度设置，不设置即只支持固定维度的参数。本例子其实可以不设置，因为我们传入的参数都是自己调整好维度的。

其实对于Bert模型的输入而言，seq_len是不需要的。可以阅读一下模型的forward()函数（前向传播的函数）定义，其实里面丢弃了seq_len。

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 "/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中)

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（展开）了，所以函数参数变多了。解决方法可以给它在套上一层。比如：

写法一

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则是整型，也因为是整型，所以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手册在这：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()方法，它的参数声明如下：

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

性能对比

模型转换为ONNX模型后，性能表现如何呢？接下来可以验证一下。先写一些辅助函数：

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的耗时应该会更低，不过这是后话了，本文暂且不表。