tensorflow入门教程(五十一)tensorflow转ONNX使用教程(Python)

#
#作者：韦访
#博客：https://blog.csdn.net/rookie_wei
#微信：1007895847
#添加微信的备注一下是CSDN的
#欢迎大家一起学习
#

---- 韦访20190813

1、概述

关于ONNX的介绍，请自行百度。现在很多深度学习框架都支持了ONNX，方便了模型的部署和在各个框架之间转移，未来可期。这一讲，就来看看怎么将tensorflow训练的模型转成ONNX以及如果调用。ONNX官网如下，

https://onnx.ai/

https://github.com/onnx/onnx

 

2、tensorflow训练模型

还是以我们写的mnist为例，不清楚的请先回去看看，链接如下，

https://blog.csdn.net/rookie_wei/article/details/80146620

这讲的例子，我们不是直接在代码里进行转换，而是用命令行的形式将tensorflow模型转成ONNX模型，这就需要将模型固化成pb文件，才能转换，相关操作参考下面博客，

https://blog.csdn.net/rookie_wei/article/details/90546290

复制第二讲底部的完整代码，保存到文件mnist_saved_model.py，进行如下修改，

将

x = tf.placeholder(tf.float32, [None, 784])

改为

x = tf.placeholder(tf.float32, [None, 784], name='input')

将

keep_prob = tf.placeholder(tf.float32)

改为

keep_prob = tf.placeholder(tf.float32, name='keep_prob')

将

y_conv = tf.nn.softmax(tf.matmul(h_fc1_drop, W_fc2) + b_fc2)

改为

y_conv = tf.nn.softmax(tf.matmul(h_fc1_drop, W_fc2) + b_fc2, name='output')

这里我们不再用tf.train.Saver()的方式保存模型，所以将

#将训练结果保存，如果不保存我们这次训练结束后的结果也随着程序运行结束而释放了

savePath = './mnist_conv/'

saveFile = savePath + 'mnist_conv.ckpt'

if os.path.exists(savePath) == False:

    os.mkdir(savePath)

 

saver = tf.train.Saver()

删除

将

# 最后，将会话保存下来

    saver.save(sess, saveFile)

改为

    # 用SavedModel的方式保存

    tf.compat.v1.saved_model.simple_save(sess,

            "./saved_model",

            inputs={"input": x, 'keep_prob':keep_prob},

            outputs={"output": y_conv})

完整代码如下，

# coding: utf-8
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
import os

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

mnist = input_data.read_data_sets('mnist_data', one_hot=True)

#初始化过滤器
def weight_variable(shape):
    return tf.Variable(tf.truncated_normal(shape, stddev=0.1))

#初始化偏置，初始化时，所有值是0.1
def bias_variable(shape):
    return tf.Variable(tf.constant(0.1, shape=shape))

#卷积运算，strides表示每一维度滑动的步长，一般strides[0]=strides[3]=1
#第四个参数可选"Same"或"VALID"，“Same”表示边距使用全0填充
def conv2d(x, W):
    return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding="SAME")

#池化运算
def max_pool_2x2(x):

    return tf.nn.max_pool(x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding="SAME")

#创建x占位符，用于临时存放MNIST图片的数据，
# [None, 784]中的None表示不限长度，而784则是一张图片的大小（28×28=784）
x = tf.placeholder(tf.float32, [None, 784], name='input')
#y_存的是实际图像的标签，即对应于每张输入图片实际的值
y_ = tf.placeholder(tf.float32, [None, 10])

#将图片从784维向量重新还原为28×28的矩阵图片,
# 原因参考卷积神经网络模型图，最后一个参数代表深度，
# 因为MNIST是黑白图片，所以深度为1,
# 第一个参数为-1,表示一维的长度不限定，这样就可以灵活设置每个batch的训练的个数了
x_image = tf.reshape(x, [-1, 28, 28, 1])

#第一层卷积
#将过滤器设置成5×5×1的矩阵，
#其中5×5表示过滤器大小，1表示深度，因为MNIST是黑白图片只有一层。所以深度为1
#32表示我们要创建32个大小5×5×1的过滤器，经过卷积后算出32个特征图（每个过滤器得到一个特征图），即输出深度为64
W_conv1 = weight_variable([5, 5, 1, 32])
#有多少个特征图就有多少个偏置
b_conv1 = bias_variable([32])
#使用conv2d函数进行卷积计算，然后再用ReLU作为激活函数
h_conv1 = tf.nn.relu(conv2d(x_image, W_conv1) + b_conv1)
#卷积以后再经过池化操作
h_pool1 = max_pool_2x2(h_conv1)

#第二层卷积
#因为经过第一层卷积运算后，输出的深度为32,所以过滤器深度和下一层输出深度也做出改变
W_conv2 = weight_variable([5, 5, 32, 64])
b_conv2 = bias_variable([64])
h_conv2 = tf.nn.relu(conv2d(h_pool1, W_conv2) + b_conv2)
h_pool2 = max_pool_2x2(h_conv2)

#全连接层
#经过两层卷积后，图片的大小为7×7（第一层池化后输出为（28/2）×（28/2），
#第二层池化后输出为（14/2）×（14/2））,深度为64，
#我们在这里加入一个有1024个神经元的全连接层，所以权重W的尺寸为[7 * 7 * 64, 1024]
W_fc1 = weight_variable([7 * 7 * 64, 1024])
#偏置的个数和权重的个数一致
b_fc1 = bias_variable([1024])
#这里将第二层池化后的张量（长：7 宽：7 深度：64） 变成向量（跟上一节的Softmax模型的输入一样了）
h_pool2_flat = tf.reshape(h_pool2, [-1, 7 * 7 * 64])
#使用ReLU激活函数
h_fc1 = tf.nn.relu(tf.matmul(h_pool2_flat, W_fc1) + b_fc1)

#dropout
#为了减少过拟合，我们在输出层之前加入dropout
keep_prob = tf.placeholder(tf.float32, name='keep_prob')
h_fc1_drop = tf.nn.dropout(h_fc1, keep_prob)

#输出层
#全连接层输入的大小为1024,而我们要得到的结果的大小是10（0～9），
# 所以这里权重W的尺寸为[1024, 10]
W_fc2 = weight_variable([1024, 10])
b_fc2 = bias_variable([10])

#最后都要经过Softmax函数将输出转化为概率问题
y_conv = tf.nn.softmax(tf.matmul(h_fc1_drop, W_fc2) + b_fc2, name='output')

#损失函数和损失优化
cross_entropy = tf.reduce_mean(-tf.reduce_sum(y_ * tf.log(y_conv)))
train_step = tf.train.AdamOptimizer(1e-4).minimize(cross_entropy)
 
#测试准确率,跟Softmax回归模型的一样
correct_prediction = tf.equal(tf.argmax(y_conv, 1), tf.argmax(y_, 1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

# #将训练结果保存，如果不保存我们这次训练结束后的结果也随着程序运行结束而释放了
# savePath = './mnist_conv/'
# saveFile = savePath + 'mnist_conv.ckpt'
# if os.path.exists(savePath) == False:
#     os.mkdir(savePath)

# saver = tf.train.Saver()

#开始训练
with tf.Session() as sess:

    #初始化所有变量
    sess.run(tf.global_variables_initializer())

    #训练两万次
    for i in range(2000):

        #每次获取50张图片数据和对应的标签
        batch = mnist.train.next_batch(50)

        #每训练100次，我们打印一次训练的准确率
        if i % 100 == 0:
            train_accuracy =sess.run(accuracy, feed_dict={x:batch[0], y_:batch[1], keep_prob:1.0})
            print("step %d, training accuracy %g" % (i, train_accuracy))

        #这里是真的训练，将数据传入
        sess.run(train_step, feed_dict={x:batch[0], y_:batch[1], keep_prob:0.5})

 

    # print ("end train, start testing...")
    # mean_value = 0.0
    # for i in range(mnist.test.labels.shape[0]):
    #     batch = mnist.test.next_batch(50)
    #     train_accuracy = sess.run(accuracy, feed_dict={x: batch[0], y_: batch[1], keep_prob: 1.0})
    #     mean_value += train_accuracy

 

    # print("test accuracy %g" % (mean_value / mnist.test.labels.shape[0]))
    # #训练结束后，我们使用mnist.test在测试最后的准确率
    # print("test accuracy %g" % sess.run(accuracy, feed_dict={x:mnist.test.images, y_:mnist.test.labels, keep_prob:1.0}))

 
    # 最后，将会话保存下来
    # saver.save(sess, saveFile)

    # 用SavedModel的方式保存
    tf.compat.v1.saved_model.simple_save(sess,
            "./saved_model",
            inputs={"input": x, 'keep_prob':keep_prob},
            outputs={"output": y_conv})

运行结果，

 在当前文件夹下新增名为saved_model的文件夹，文件夹内容如下，

3、调用训练好的模型

模型训练好了以后，最好写个代码来调用它，以验证我们得到的模型是正确的。我这里就直接贴完整的代码了，代码如下，

# coding: utf-8

import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
import os
import numpy as np

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

mnist = input_data.read_data_sets('mnist_data', one_hot=True)

with tf.Session() as sess:
    
    tf.compat.v1.saved_model.loader.load (sess, ["serve"], "./saved_model")

    x = sess.graph.get_tensor_by_name('input:0')
    y = sess.graph.get_tensor_by_name('output:0')
    keep_prob = sess.graph.get_tensor_by_name('keep_prob:0')

    for i in range(mnist.test.labels.shape[0]):

        batch = mnist.test.next_batch(1)
        print("wf======== batch :",  np.argmax(batch[1], 1))
        train_accuracy = sess.run(y, feed_dict={x: batch[0], keep_prob:1})
        print("wf======== accuracy :",  np.argmax(train_accuracy, 1))
        if i == 10:
            break

运行结果，

wf======== batch : [7]

wf======== accuracy : [7]

wf======== batch : [8]

wf======== accuracy : [8]

wf======== batch : [1]

wf======== accuracy : [1]

wf======== batch : [7]

wf======== accuracy : [7]

wf======== batch : [6]

wf======== accuracy : [6]

wf======== batch : [2]

wf======== accuracy : [2]

wf======== batch : [6]

wf======== accuracy : [6]

wf======== batch : [7]

wf======== accuracy : [7]

wf======== batch : [6]

wf======== accuracy : [6]

wf======== batch : [1]

wf======== accuracy : [1]

wf======== batch : [8]

wf======== accuracy : [8]

说明我们的模型调用没问题。

4、将tensorflow模型转换成ONNX模型

先安装tf2onnx模块，执行下面命令，

pip install -U tf2onnx

接着，就要将上面的模型转成ONNX的模型，用一个命令行就可以搞定了，命令如下，

python -m tf2onnx.convert   --saved-model ./saved_model   --output ./mnist_onnx/mnist.onnx   --opset 7

运行结果，

2019-08-13 16:42:58,593 - INFO - Using tensorflow=1.13.1, onnx=1.4.1, tf2onnx=1.5.3/7b598d

2019-08-13 16:42:58,593 - INFO - Using opset

2019-08-13 16:42:58,595 - WARNING - Shape of placeholder keep_prob is unknown, treated it as a scalar

2019-08-13 16:42:58,888 - INFO - Optimizing ONNX model

2019-08-13 16:42:59,325 - INFO - After optimization: Add -2 (2->0), Cast -2 (2->0), Const +1 (10->11), Identity -1 (1->0), Reshape +1 (2->3), Transpose -7 (8->1)

2019-08-13 16:42:59,368 - INFO -

2019-08-13 16:42:59,368 - INFO - Successfully converted TensorFlow model ./saved_model to ONNX

2019-08-13 16:42:59,387 - INFO - ONNX model is saved at ./mnist_onnx/mnist.onnx

当前文件夹下生成一个新的名为mnist_onnx的文件夹，该文件夹下有一个mnist.onnx文件，这个就是我们转换后的ONNX模型文件，详情可以参考一下链接，

https://github.com/onnx/tutorials/blob/master/tutorials/TensorflowToOnnx-1.ipynb

想知道更多其他深度学习模型的转换，参考下面链接，

https://github.com/onnx/tutorials

5、使用ONNX调用转换后的模型

得到模型后，我们最终的目的是用起来，我偷了个懒，mnist数据的获取我还是通过了tensorflow给的方法，但是在预测部分，则是用ONNX的方法，比较简单，直接给代码，如下，

import onnxruntime as ort
import numpy as np
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data

session = ort.InferenceSession("./mnist_onnx/mnist.onnx", None)
keep_prob = session.get_inputs()[0].name
x = session.get_inputs()[1].name

y = session.get_outputs()[0].name

print("wf========keep_prob name:", keep_prob)
print("wf========input name:", x)
print("wf========output name:", y)

mnist = input_data.read_data_sets('mnist_data', one_hot=True)

for i in range(mnist.test.labels.shape[0]):
    batch = mnist.test.next_batch(1)
    print("wf======== batch :",  np.argmax(batch[1], 1))
    train_accuracy = session.run(None, {x: batch[0],  keep_prob:np.ones_like(batch[0])})
    print("wf======== accuracy :",  np.argmax(train_accuracy))
    if i == 10:
        break

运行结果，

wf======== batch : [8]

wf======== accuracy : 8

wf======== batch : [6]

wf======== accuracy : 6

wf======== batch : [1]

wf======== accuracy : 1

wf======== batch : [8]

wf======== accuracy : 2

wf======== batch : [3]

wf======== accuracy : 3

wf======== batch : [2]

wf======== accuracy : 2

wf======== batch : [0]

wf======== accuracy : 0

wf======== batch : [5]

wf======== accuracy : 5

wf======== batch : [0]

wf======== accuracy : 0

wf======== batch : [8]

wf======== accuracy : 8

wf======== batch : [3]

wf======== accuracy : 3

 

后话：

ONNX是微软等公司主导的开源项目，在Windows平台上部署是很方便的，支持的语言也很多，而且很多深度学习框架都支持ONNX，所以是个值得注意的东西。