[ONNX] Train in Tensorflow and export to ONNX (Part II)

If you read the previous post as the link below, you probably may ask a question: If the input TF graph for freezing is not a binary format, what do we do?
http://danny270degree.blogspot.com/2018/08/onnx-train-in-tensorflow-and-export-to.html

Let us recall the previous example below. The file "graph.proto" is the binary format of the protobuf file for TensorFlow graph generated from the following function:

  with open("graph.proto", "wb") as file:
    graph = tf.get_default_graph().as_graph_def(add_shapes=True)
    file.write(graph.SerializeToString())

How to generate text format of protobuf file?
Here you go:

tf.train.write_graph(sess.graph_def, './my_mnist', 'graph.pbtxt', as_text=True)

Now, I will use my example to convert TensorFlow's model to ONNX model by myself again.
The following steps are a little bit different from the previous post. You can take both into considering your solution.

First, I train my a simple CNN model with MNIST dataset and it will save the result to ./my_mnist folder.
A simple CNN model:

# To support both python 2 and python 3
from __future__ import division, print_function, unicode_literals

# Common imports
import numpy as np
import tensorflow as tf
import os

def reset_graph(seed=42):
    tf.reset_default_graph()
    tf.set_random_seed(seed)
    np.random.seed(seed)


n_epochs = 1
batch_size = 10
height = 28
width = 28
channels = 1
n_inputs = height * width

conv1_fmaps = 32
conv1_ksize = 3
conv1_stride = 1
conv1_pad = "SAME"

conv2_fmaps = 64
conv2_ksize = 3
conv2_stride = 2
conv2_pad = "SAME"

pool3_fmaps = conv2_fmaps

n_fc1 = 64
n_outputs = 10

reset_graph()

with tf.name_scope("inputs"):
    X = tf.placeholder(tf.float32, shape=[None, n_inputs], name="X")
    X_reshaped = tf.reshape(X, shape=[-1, height, width, channels])
    y = tf.placeholder(tf.int32, shape=[None], name="y")

conv1 = tf.layers.conv2d(X_reshaped, filters=conv1_fmaps, kernel_size=conv1_ksize,
                         strides=conv1_stride, padding=conv1_pad,
                         activation=tf.nn.relu, name="conv1")
conv2 = tf.layers.conv2d(conv1, filters=conv2_fmaps, kernel_size=conv2_ksize,
                         strides=conv2_stride, padding=conv2_pad,
                         activation=tf.nn.relu, name="conv2")

with tf.name_scope("pool3"):
    pool3 = tf.nn.max_pool(conv2, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding="VALID")
    pool3_flat = tf.reshape(pool3, shape=[-1, pool3_fmaps * 7 * 7])

with tf.name_scope("fc1"):
    fc1 = tf.layers.dense(pool3_flat, n_fc1, activation=tf.nn.relu, name="fc1")

with tf.name_scope("output"):
    logits = tf.layers.dense(fc1, n_outputs, name="output")
    Y_proba = tf.nn.softmax(logits, name="Y_proba")

with tf.name_scope("train"):
    xentropy = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=y)
    loss = tf.reduce_mean(xentropy)
    optimizer = tf.train.AdamOptimizer()
    training_op = optimizer.minimize(loss)

with tf.name_scope("eval"):
    correct = tf.nn.in_top_k(logits, y, 1)
    accuracy = tf.reduce_mean(tf.cast(correct, tf.float32))

with tf.name_scope("init_and_save"):
    init = tf.global_variables_initializer()
    saver = tf.train.Saver()

from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets("/tmp/MNIST_data/data/")

with tf.Session() as sess:
    tf.train.write_graph(sess.graph_def, './my_mnist', 'graph.pbtxt', as_text=True) 
    init.run()
    for epoch in range(n_epochs):
        for iteration in range(mnist.train.num_examples // batch_size):
            X_batch, y_batch = mnist.train.next_batch(batch_size)
            sess.run(training_op, feed_dict={X: X_batch, y: y_batch})
            print("iteration", iteration)
        #acc_train = accuracy.eval(feed_dict={X: X_batch, y: y_batch})
        #acc_test = accuracy.eval(feed_dict={X: mnist.test.images, y: mnist.test.labels})
        #print(epoch, "Train accuracy:", acc_train, "Test accuracy:", acc_test)

        save_path = saver.save(sess, "./my_mnist/my_mnist_model")

my_mnist/

├── checkpoint

├── graph.pbtxt

├── my_mnist_model.data-00000-of-00001

├── my_mnist_model.index

└── my_mnist_model.meta

Second, freeze the model by TensorFlow's freeze_graph:  ( the format of input graph is text format )

python -m tensorflow.python.tools.freeze_graph --input_graph=./graph.pbtxt \
--input_checkpoint=./my_mnist_model \
--input_binary=false \
--output_graph=./frozen_graph.pb \
--output_node_names=output/output/BiasAdd

P.S: you need to find out your output_node_names, basically, it usually will be the op after logit function.

Third, convert the model to ONNX format:
Here I use tensorflow-onnx to convert model and will get "model.onnx" file.

python3 -m tf2onnx.convert \
    --input frozen_graph.pb \
    --inputs inputs/X:0[1,784] \
    --outputs output/output/BiasAdd:0 \
    --output model.onnx \
    --verbose                               

P.S: You need to use Python3 to do converting because it is known issue for Ptyhon2 to execute tf2onnx: https://github.com/onnx/tensorflow-onnx/issues/60
P.S: the code in convert.py contains the way to change/determine tensors' shape.

Now, I use Netron to visualize my model in graph view on the web. Here it looks like:

If I dump the file: model.onnx, and it looks like the following code:

graph tf2onnx{
  initializers: {
    tensor <2 > %pool3/Reshape/shape:0,
    tensor <64 10 > %output/kernel:0,
    tensor <64 > %conv2/bias:0,
    tensor <10 > %output/bias:0,
    tensor <32 1 3 3 > %conv1/kernel:0,
    tensor <4 > %inputs/Reshape/shape:0,
    tensor <64 > %fc1/bias:0,
    tensor <64 32 3 3 > %conv2/kernel:0,
    tensor <32 > %conv1/bias:0,
    tensor <3136 64 > %fc1/kernel:0
  },
  inputs : {
    INT64 tensor <2> %pool3/Reshape/shape:0,
    FLOAT tensor <64, 10> %output/kernel:0,
    FLOAT tensor <64> %conv2/bias:0,
    FLOAT tensor <10> %output/bias:0,
    FLOAT tensor <32, 1, 3, 3> %conv1/kernel:0,
    INT64 tensor <4> %inputs/Reshape/shape:0,
    FLOAT tensor <64> %fc1/bias:0,
    FLOAT tensor <64, 32, 3, 3> %conv2/kernel:0,
    FLOAT tensor <32> %conv1/bias:0,
    FLOAT tensor <3136, 64> %fc1/kernel:0,
    FLOAT tensor <1, 784> %inputs/X:0
  }
}
[inputs/Reshape] UNDEFINED tensor <> %inputs/Reshape:0 = Reshape(FLOAT tensor <1, 784> %inputs/X:0, INT64 tensor <4> %inputs/Reshape/shape:0)
[conv1/Conv2D__2] UNDEFINED tensor <> %conv1/Conv2D__2:0 = Transpose <perm:INTS [0,3,1,2]> (UNDEFINED tensor <> %inputs/Reshape:0)
[conv1/Conv2D] UNDEFINED tensor <> %conv1/Conv2D:0 = Conv <kernel_shape:INTS [3,3]> , pads:INTS [1,1,1,1]> , dilations:INTS [1,1]> , strides:INTS [1,1]> (UNDEFINED tensor <> %conv1/Conv2D__2:0, FLOAT tensor <32, 1, 3, 3> %conv1/kernel:0)
[conv1/Conv2D__3] UNDEFINED tensor <> %conv1/Conv2D__3:0 = Transpose <perm:INTS [0,2,3,1]> (UNDEFINED tensor <> %conv1/Conv2D:0)
[conv1/BiasAdd] UNDEFINED tensor <> %conv1/BiasAdd:0 = Add(UNDEFINED tensor <> %conv1/Conv2D__3:0, FLOAT tensor <32> %conv1/bias:0)
[conv1/Relu] UNDEFINED tensor <> %conv1/Relu:0 = Relu(UNDEFINED tensor <> %conv1/BiasAdd:0)
[conv2/Conv2D__4] UNDEFINED tensor <> %conv2/Conv2D__4:0 = Transpose <perm:INTS [0,3,1,2]> (UNDEFINED tensor <> %conv1/Relu:0)
[conv2/Conv2D] UNDEFINED tensor <> %conv2/Conv2D:0 = Conv <kernel_shape:INTS [3,3]> , pads:INTS [0,0,1,1]> , dilations:INTS [1,1]> , strides:INTS [2,2]> (UNDEFINED tensor <> %conv2/Conv2D__4:0, FLOAT tensor <64, 32, 3, 3> %conv2/kernel:0)
[conv2/Conv2D__5] UNDEFINED tensor <> %conv2/Conv2D__5:0 = Transpose <perm:INTS [0,2,3,1]> (UNDEFINED tensor <> %conv2/Conv2D:0)
[conv2/BiasAdd] UNDEFINED tensor <> %conv2/BiasAdd:0 = Add(UNDEFINED tensor <> %conv2/Conv2D__5:0, FLOAT tensor <64> %conv2/bias:0)
[conv2/Relu] UNDEFINED tensor <> %conv2/Relu:0 = Relu(UNDEFINED tensor <> %conv2/BiasAdd:0)
[pool3/MaxPool__6] UNDEFINED tensor <> %pool3/MaxPool__6:0 = Transpose <perm:INTS [0,3,1,2]> (UNDEFINED tensor <> %conv2/Relu:0)
[pool3/MaxPool] UNDEFINED tensor <> %pool3/MaxPool:0 = MaxPool <strides:INTS [2,2]> , kernel_shape:INTS [2,2]> (UNDEFINED tensor <> %pool3/MaxPool__6:0)
[pool3/MaxPool__7] UNDEFINED tensor <> %pool3/MaxPool__7:0 = Transpose <perm:INTS [0,2,3,1]> (UNDEFINED tensor <> %pool3/MaxPool:0)
[pool3/Reshape] UNDEFINED tensor <> %pool3/Reshape:0 = Reshape(UNDEFINED tensor <> %pool3/MaxPool__7:0, INT64 tensor <2> %pool3/Reshape/shape:0)
[fc1/fc1/MatMul] UNDEFINED tensor <> %fc1/fc1/MatMul:0 = MatMul(UNDEFINED tensor <> %pool3/Reshape:0, FLOAT tensor <3136, 64> %fc1/kernel:0)
[fc1/fc1/BiasAdd] UNDEFINED tensor <> %fc1/fc1/BiasAdd:0 = Add(UNDEFINED tensor <> %fc1/fc1/MatMul:0, FLOAT tensor <64> %fc1/bias:0)
[fc1/fc1/Relu] UNDEFINED tensor <> %fc1/fc1/Relu:0 = Relu(UNDEFINED tensor <> %fc1/fc1/BiasAdd:0)
[output/output/MatMul] UNDEFINED tensor <> %output/output/MatMul:0 = MatMul(UNDEFINED tensor <> %fc1/fc1/Relu:0, FLOAT tensor <64, 10> %output/kernel:0)
[output/output/BiasAdd] FLOAT tensor <1, 10> %output/output/BiasAdd:0 = Add(UNDEFINED tensor <> %output/output/MatMul:0, FLOAT tensor <10> %output/bias:0)
  return FLOAT tensor <1, 10> %output/output/BiasAdd:0