Swap a TensorFlow Dataset input pipeline with a placeholder after training

Question

I'm working with the new tf.data.Dataset API and I can't seem to figure out how to perform inference. Ultimately, I want to convert my model to a TensorRT graph and run it on the TX2, and all of the examples I have found assume you have a tf.placeholder for the input. Here is pseudocode for how I am training. The [...] is just meant to be a placeholder since I didn't actually run the code. Let's not debate the model, as it is just suppose to give an example:

import tensorflow as tf

# Setup iterator
datain = tf.data.FixedLengthRecordDataset(datafiles, record_bytes1)
labels = tf.data.FixedLengthRecordDataset(labelfiles, record_bytes2)
dataset = tf.data.Dataset.zip((datain, labels))
dataset = dataset.prefetch(batch_size)
dataset = dataset.repeat(n_epoch)
iterator = dataset.make_initializable_iterator()

sess = tf.Session()
sess.run(iterator.initializer)
[batch_x, batch_y] = iterator.get_next()

# Define model function (let's not debate model except as relevant to question)
def model_fn(xin):
    x0 = tf.transpose(tf.reshape(xin, [...], name='input'))
    w = tf.Variable(tf.truncated_normal([...], stddev=0.1))
    x1 = tf.nn.conv2d(x0, w, strides=[...], padding='VALID')
    b = tf.Variable(tf.constant(0.0, shape=[...]))
    x2 = tf.nn.bias_add(x1, b)
    x3 = tf.nn.relu(x2, name='output')
    return x3

# Setup training environment
model = model_fn(batch_x)
loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(logits=model, labels=batch_y))
optimizer = tf.train.AdamOptimizer(learning_rate=1e-3).minimize(loss)

# Train Model
while True:
    try:
        sess.run(optimizer)
    except tf.errors.OutOfRangeError:
        break

# Save model
saver = tf.train.Saver(name='saver')
saver.save(sess, 'temp/path')

My question is how do I get this into TensorRT without having the input be a tf.placeholder? All of the example I can find use a tf.placeholder as the input. This example suggests that I can replace the iterator with a placeholder using the SavedModel class, but I cannot seem to find any documentation on how to accomplish that.

Thanks!

EDIT: Here is my solution thanks to the help below

from tensorflow.python.tools import optimize_for_inference_lib
import uff

# You can feed data to the IteratorGetNext node using feed_dict
input_node_name = 'iterator_scope_name/IteratorGetNext'
output_node_name = 'model_scope_name/output'

# Run inference on the trained model:
graph = tf.get_default_graph()
batch_x = graph.get_tensor_by_name(input_node_name + ':0')
networkout = graph.get_tensor_by_name(output_node_name + ':0')
testdata, testlabel = custom_data_reader_fn(data_folder)
# This will evaluate the model
label = sess.run(networkout, feed_dict={batch_x: testdata})

# Freeze model and create a UFF file:
graph_def = graph.as_graph_def() # Convert the graph to a serialized pb
frozen_graph_def = tf.graph_util.convert_variables_to_constants(sess,
    graph_def, [output_node_name])
opt_graph_def = optimize_for_inference_lib.optimize_for_inference(
    frozen_graph_def, [input_node_name], [output_node_name],
    tf.float32.as_datatype_enum)
uff.from_tensorflow(opt_graph_def, [output_node_name], quiet=False,
    output_filename='opt_model.uff')

that will write out a UFF file that TensorRT can utilize. The biggest issues that I encountered was:

1. I didn't realize that the optimize_for_inference_lib.optimize_for_inference operation replaced the iterator with a tf.placeholder
2. I did not know what node to feed data to for evaluation: you can feed data to the IteratorGetNext node

Saari

Are you constrained to use this relatively low-level training scheme or would you consider a solution using tf.estimator for training and exporting also viable? – Embroidery

if tf.estimator will work with TensorRT, I am happy to use that. I'm not sure it will though. – Saari

As far as I understand it, TensorRT works with trained graphs. Estimator helps you with the training and exporting of the graph, then you freeze it and import it in TensorRT. I'll be able to answer properly only tomorrow, I'll try to explain how to set up the various pieces of code you need for the training and exporting – Embroidery

Hi GPhilo. I have a trained graph, but it takes in an iterator. I believe that if I could just understand how to feed the graph with new data, or replace the iterator with a tf.placeholder after training, I would be set. I am interested to see how to set it up using the estimator though. – Saari

Ok, if all you need is that, have a look at this: github.com/tensorflow/tensorflow/blob/master/tensorflow/python/… you'll need to pass on the graph.txt file that should be in your checkpoint folder, make sure to read the comments at the top about using the non-frozen graph – Embroidery

Have a look at this question in the opposite direction (placeholder -> iterator). The answer is basically the same. – Ouse

Answer 1

Since you already have a trained graph saved in a checkpoint, in theory the simplest solution for you is to export the inference graph via optimize_for_inference.

This tool works both for already-frozen graphs and, as is your case, for graphs with variables still defined. Assuming you go for the frozen graph way, the first step is to transform your graph's variables in constants via:

python freeze_graph.py \
--input_graph=temp/path/graph.pbtxt \
--input_checkpoint=temp/path/your_model_name.ckpt \
--output_graph=frozen_model.pb \
--output_node_names=name_of_the_output_tensor_you_want_to_use

This will generate a new binary file called frozen_model.pb that has the Variable operations replaced with Const ops with the values loaded from the checkpoint file.

Then, you need to generate the inference graph with:

python optimize_for_inference.py \
--input=frozen_model.pb \
--output=inference.pb \
--frozen_graph=True \
--input_names=IteratorGetNext
--output_names=name_of_the_output_tensor_you_want_to_use

This will replace the IteratorGetNext node with a float placeholder. You might want to choose another node, in which case just change the name. You can also change the type of the generated placeholder via the --placeholder_type_enum option. In that case, you need to provide an integer value matching the datatype you want from the DataType enum.

NOTE: I said "in theory" because actually inspecting the generated inception graph from a test I made it seems there are still some weird ops in there that are not really necessary for inference. You might have to further process your graph via nvidia's Graph Surgeon or TF's graph transform tool

Embroidery

Thanks for the answer. My biggest hickup was not knowing: 1) what TensorFlow node to feed data to (here it is IteratorGetNext) and 2) not realizing that the optimize_for_inference step actually replaced the IteratorGetNext with a placeholder. – Saari

Yeah, for all the examples and documentation there is out there, I feel the step between training and actually using the network is still pretty badly documented. There's at least three things to do this and most of them require a bunch of non-obvious steps... I wish they could simplify this a little bit :/ – Embroidery

Agreed. I edited my question giving the code for the inference and creation of the UFF file for TensorRT. Thanks for your help! – Saari

@Embroidery I tried the optimize_for_inference but it did not replace the IteratorGetNext node with a placeholder. Am I missing some flag ? I mean, there is a new Placeholder but seems like the dataset loader nodes are still present in the graph and when I try to convert the model using a compiler frontend, it throws exception on ops IteratorGetNext and IteratorV2 – Pianette

What version of TF are you using? optimize_for_inference is becoming a bit outdated and it's possible it might not work with the new modifications introduced (I saw a GitHub issue where they recommend using transform_graph instead) – Embroidery