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

书接上回

在之前的文章《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 "/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存储的便是参数输入参数。 其实很多人都遇到过这个问题，比如：

github.com/pytorch/pyt…

github.com/onnx/onnx/i…

应该是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模型完结一个在线猜测服务。但限于文章篇幅，咱们下次再聊！请我们继续重视。