I am using TensorRT 4.0.1.6 with Tensorflow, and realized that,
if I use “tf.reshape”, TensorRT gives an error like below.
ERROR: UFFParser: Parser error: reshape1: Reshape: Volume mismatch
I also found that this issue was already known: https://devtalk.nvidia.com/default/topic/1036882/jetson-tx2/volume-mismatch-error-at-tensorrt-engine/post/5269001/#5269001
Are there ways to work around this?
Should I replace the part with a custom plugin layer?
For information, the below is the minimal example code for this.
(Although tf.reshape() is the last layer in this example, in my actual application, tf.reshape() is used in the middle of network multiple times.)
Frozen Graph
""" minial example of exporting network to frozen graph """
with tf.Graph().as_default():
""" some network """
image = tf.placeholder(tf.float32, [1, 768, 2560, 3])
net = slim.conv2d(image, 4, [3, 3], scope='conv1')
net = slim.max_pool2d(net, [16, 16], scope='pool1')
net = slim.conv2d(net, 4, [3, 3], scope='conv2')
net = slim.max_pool2d(net, [2, 2], scope='pool2')
net = tf.reshape(net, (-1, 2), name='reshape1')
gpu_config = tf.ConfigProto(allow_soft_placement=True)
gpu_config.gpu_options.allow_growth = True
with tf.Session(config=gpu_config) as sess:
init_op = tf.global_variables_initializer()
sess.run(init_op)
""" specify tensors that I will need when doing inference """
output_tensors = [net]
output_names = [n.op.name for n in output_tensors]
print("output tensors: ", output_names)
graphdef = tf.get_default_graph().as_graph_def()
frozen_graph = tf.graph_util.convert_variables_to_constants(sess, graphdef, output_names)
frozen_graph = tf.graph_util.remove_training_nodes(frozen_graph)
tf.train.write_graph(frozen_graph, '.', 'reshape-test.pb', as_text=False)
UFF Conversion
$ convert-to-uff tensorflow --input-file reshape-test.pb -O reshape1
Loading reshape-test.pb
Using output node reshape1
Converting to UFF graph
No. nodes: 16