Hi All,
I am new to tensorRT. I want to convert a tensorflow UFF model to tensorrt PLAN file on Jetson Tx2 using TRT 4.1. First I try to convert Tensorflow Frozen Graph to UFF model which gives me a
warning : No conversion function registered for layer: BatchMatMul yet.
Is it really true that tensorrt doesn’t have BatchMatMul OP implemented ?
Here is the example tensorflow code I used:
import tensorflow as tf
from tensorflow.python.tools import freeze_graph
import numpy as np
batch_size, n, m, k = 10, 3, 3, 3
#A = tf.Variable(tf.ones(shape=(batch_size, n, m)))
D = tf.Variable(tf.ones(shape=(10, 3,3)))
A = tf.placeholder(tf.float32, shape=(batch_size, n, m))
B = tf.placeholder(tf.float32, shape=(batch_size, m, k))
B = 2*B
A = tf.multiply(D,A)
C = tf.matmul(A, B, name='outputnode')
with tf.Session() as sess:
A_rand = np.random.rand(batch_size, n, m)
B_rand = np.random.rand(batch_size, m, k)
initialize = tf.global_variables_initializer()
sess.run(initialize, feed_dict={A: A_rand, B: B_rand})
tf.train.write_graph(sess.graph.as_graph_def(),'.','./batchmatmul.pbtxt', as_text=True)
saver = tf.train.Saver()
saver.save(sess, './batchmatmul.ckpt')
train_writer = tf.summary.FileWriter("./logs")
train_writer.add_graph(sess.graph)
freeze_graph.freeze_graph('./batchmatmul.pbtxt', "",False,'./batchmatmul.ckpt', "outputnode", "save/restore_all", "save/Const:0",'./batchmatmul.pb', True,"")
import uff
with open('./batchmatmul.pb', 'rb') as f:
frozen_graph.ParseFromString(f.read())
uff.from_tensorflow(graphdef=frozen_graph,
output_filename="test.uff",
output_nodes="outputnode",
text=True)
Assuming that TensorRT doesn’t have BatchMatMul OP, should I write my own plugin in c++ to do batch matrix multiplication ?
Any help would be appreciated.
Thank you!