模型部署之 ONNX ONNXRuntime

通常我们在训练模型时可以使用很多不同的框架,比如有的同学喜欢用 Pytorch,有的同学喜欢使用 TensorFLow,也有的喜欢 MXNet,以及深度学习最开始流行的 Caffe等等,这样不同的训练框架就导致了产生不同的模型结果包,在模型进行部署推理时就需要不同的依赖库,而且同一个框架比如tensorflow 不同的版本之间的差异较大, 为了解决这个混乱问题,LF AI 这个组织联合 Facebook, MicroSoft等公司制定了机器学习模型的标准,这个标准叫做ONNX, Open Neural Network Exchage,所有其他框架产生的模型包 (.pth, .pb) 都可以转换成这个标准格式,转换成这个标准格式后,就可以使用统一的 ONNX Runtime等工具进行统一部署。

这其实可以和 JVM 对比,
A Java virtual machine (JVM) is a virtual machine that enables a computer to run Java programs as well as programs written in other languages that are also compiled to Java bytecode. The JVM is detailed by a specification that formally describes what is required in a JVM implementation. Having a specification ensures interoperability of Java programs across different implementations so that program authors using the Java Development Kit (JDK) need not worry about idiosyncrasies of the underlying hardware platform.

JAVA中有 JAVA 语言 + .jar 包 + JVM,同时还有其他的语言比如 Scala等也是建立在 JVM上运行的,因此不同的语言只要都最后将程序转换成 JVM可以统一识别的格式,就可以在统一的跨平台 JVM JAVA 虚拟机上运行。这里JVM使用的 包是二进制包,因此里面的内容是不可知的,人类难以直观理解的。

这里 ONNX 标准采取了谷歌开发 protocal buffers 作为格式标准,这个格式是在 XML, json的基础上发展的,是一个人类易理解的格式。ONNX 官网对ONNX的介绍如下:
ONNX defines a common set of operators - the building blocks of machine learning and deep learning models - and a common file format to enable AI developers to use models with a variety of frameworks, tools, runtimes, and compilers.
ONNX支持的模型来源,基本上囊括了我们日常使用的所有框架:

ONNX支持的模型来源

ONNX的文件格式,采用的是谷歌的 protocal buffers,和 caffe采用的一致。


ONNX定义的数据类包括了我们常用的数据类型,用来定义模型中的输出输出格式

ONNX中定义了很多我们常用的节点,比如 Conv,ReLU,BN, maxpool等等约124种,同时也在不停地更新中,当遇到自带节点库中没有的节点时,我们也可以自己写一个节点

  • 有了输入输出,以及计算节点,就可以根据 pytorch框架中的 forward 记录一张模型从输入图片到输出的计算图,ONNX 就是将这张计算图用标准的格式存储下来了,可以通过一个工具 Netron对 ONNX 进行可视化,如第一张图右侧所示;
  • 保存成统一的 ONNX 格式后,就可以使用统一的运行平台来进行 inference。

pytorch原生支持 ONNX 格式转码,下面是实例:

1. 将pytorch模型转换为onnx格式,直接傻瓜式调用 torch.onnx.export(model, input, output_name)

import torch
from torchvision import models

net = models.resnet.resnet18(pretrained=True)
dummpy_input = torch.randn(1,3,224,224)
torch.onnx.export(net, dummpy_input, 'resnet18.onnx')

2. 对生成的 onnx 进行查看

import onnx

# Load the ONNX model
model = onnx.load("resnet18.onnx")

# Check that the IR is well formed
onnx.checker.check_model(model)

# Print a human readable representation of the graph
print(onnx.helper.printable_graph(model.graph))

输出:
可以看到其输出有3个dict,一个是 input, 一个是 initializers,以及最后一个是operators把输入和权重 initialization 进行类似于 forward操作,在最后一个dict operators中其返回是 %191,也就是 gemm 最后一个全连接的输出

graph torch-jit-export (
  %input.1[FLOAT, 1x3x224x224]
) initializers (
  %193[FLOAT, 64x3x7x7]
  %194[FLOAT, 64]
  %196[FLOAT, 64x64x3x3]
  %197[FLOAT, 64]
  %199[FLOAT, 64x64x3x3]
  %200[FLOAT, 64]
  %202[FLOAT, 64x64x3x3]
  %203[FLOAT, 64]
  %205[FLOAT, 64x64x3x3]
  %206[FLOAT, 64]
  %208[FLOAT, 128x64x3x3]
  %209[FLOAT, 128]
  %211[FLOAT, 128x128x3x3]
  %212[FLOAT, 128]
  %214[FLOAT, 128x64x1x1]
  %215[FLOAT, 128]
  %217[FLOAT, 128x128x3x3]
  %218[FLOAT, 128]
  %220[FLOAT, 128x128x3x3]
  %221[FLOAT, 128]
  %223[FLOAT, 256x128x3x3]
  %224[FLOAT, 256]
  %226[FLOAT, 256x256x3x3]
  %227[FLOAT, 256]
  %229[FLOAT, 256x128x1x1]
  %230[FLOAT, 256]
  %232[FLOAT, 256x256x3x3]
  %233[FLOAT, 256]
  %235[FLOAT, 256x256x3x3]
  %236[FLOAT, 256]
  %238[FLOAT, 512x256x3x3]
  %239[FLOAT, 512]
  %241[FLOAT, 512x512x3x3]
  %242[FLOAT, 512]
  %244[FLOAT, 512x256x1x1]
  %245[FLOAT, 512]
  %247[FLOAT, 512x512x3x3]
  %248[FLOAT, 512]
  %250[FLOAT, 512x512x3x3]
  %251[FLOAT, 512]
  %fc.bias[FLOAT, 1000]
  %fc.weight[FLOAT, 1000x512]
) {
  %192 = Conv[dilations = [1, 1], group = 1, kernel_shape = [7, 7], pads = [3, 3, 3, 3], strides = [2, 2]](%input.1, %193, %194)
  %125 = Relu(%192)
  %126 = MaxPool[kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [2, 2]](%125)
  %195 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%126, %196, %197)
  %129 = Relu(%195)
  %198 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%129, %199, %200)
  %132 = Add(%198, %126)
  %133 = Relu(%132)
  %201 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%133, %202, %203)
  %136 = Relu(%201)
  %204 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%136, %205, %206)
  %139 = Add(%204, %133)
  %140 = Relu(%139)
  %207 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [2, 2]](%140, %208, %209)
  %143 = Relu(%207)
  %210 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%143, %211, %212)
  %213 = Conv[dilations = [1, 1], group = 1, kernel_shape = [1, 1], pads = [0, 0, 0, 0], strides = [2, 2]](%140, %214, %215)
  %148 = Add(%210, %213)
  %149 = Relu(%148)
  %216 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%149, %217, %218)
  %152 = Relu(%216)
  %219 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%152, %220, %221)
  %155 = Add(%219, %149)
  %156 = Relu(%155)
  %222 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [2, 2]](%156, %223, %224)
  %159 = Relu(%222)
  %225 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%159, %226, %227)
  %228 = Conv[dilations = [1, 1], group = 1, kernel_shape = [1, 1], pads = [0, 0, 0, 0], strides = [2, 2]](%156, %229, %230)
  %164 = Add(%225, %228)
  %165 = Relu(%164)
  %231 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%165, %232, %233)
  %168 = Relu(%231)
  %234 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%168, %235, %236)
  %171 = Add(%234, %165)
  %172 = Relu(%171)
  %237 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [2, 2]](%172, %238, %239)
  %175 = Relu(%237)
  %240 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%175, %241, %242)
  %243 = Conv[dilations = [1, 1], group = 1, kernel_shape = [1, 1], pads = [0, 0, 0, 0], strides = [2, 2]](%172, %244, %245)
  %180 = Add(%240, %243)
  %181 = Relu(%180)
  %246 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%181, %247, %248)
  %184 = Relu(%246)
  %249 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [1, 1, 1, 1], strides = [1, 1]](%184, %250, %251)
  %187 = Add(%249, %181)
  %188 = Relu(%187)
  %189 = GlobalAveragePool(%188)
  %190 = Flatten[axis = 1](%189)
  %191 = Gemm[alpha = 1, beta = 1, transB = 1](%190, %fc.weight, %fc.bias)
  return %191
}

3. 对生成的ONNX进行可视化:

onnx的可是支持有两个,一个是 netron, 一个是百度飞桨开发的visualDL
这里介绍 netron的下载安装:https://github.com/lutzroeder/Netron,对于 mac用户可以安装成功直接打开软件进行图形化选取onnx地址就可以打开

||
netron可视化图

4. ONNX Runtime

支持ONNX的runtime就是类似于JVM将统一的ONNX格式的模型包运行起来,包括对ONNX 模型进行解读,优化(融合conv-bn等操作),运行。

支持ONNX格式的runtime

这里介绍 microsoft 开发的 ONNX Runtime

4.1 ONNXRuntime的安装

https://github.com/microsoft/onnxruntime
对于使用cpu来进行推理的 mac os 可以使用

brew install libomp
pip install onnxruntime

推理

import onnxruntime as rt
import numpy as  np
data = np.array(np.random.randn(1,3,224,224))
sess = rt.InferenceSession('resnet18.onnx')
input_name = sess.get_inputs()[0].name
label_name = sess.get_outputs()[0].name

pred_onx = sess.run([label_name], {input_name:data.astype(np.float32)})[0]
print(pred_onx)
print(np.argmax(pred_onx)

可以看到,这样推理就不需要其他各种各样的pytorch等依赖,方便部署。

推荐两个易懂的视频讲解:
Everything You Want to Know About ONNX
MicroSoft onnx and onnx runtim

©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 158,847评论 4 362
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 67,208评论 1 292
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 108,587评论 0 243
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 43,942评论 0 205
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 52,332评论 3 287
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 40,587评论 1 218
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 31,853评论 2 312
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 30,568评论 0 198
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 34,273评论 1 242
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 30,542评论 2 246
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 32,033评论 1 260
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 28,373评论 2 253
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 33,031评论 3 236
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 26,073评论 0 8
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 26,830评论 0 195
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 35,628评论 2 274
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 35,537评论 2 269

推荐阅读更多精彩内容