训练后用占位符交换 TensorFlow 数据集输入管道

Question

我正在使用新的 tf.data.Dataset API，但我似乎无法弄清楚如何执行推理。最终，我想将我的模型转换为 TensorRT 图并运行它在 TX2 上，并且所有 examples I have found 假设您有一个 tf.placeholder 作为输入。这是我如何训练的伪代码。 [...] 只是一个占位符，因为我实际上没有运行代码。不讨论模型，只是想举个例子：

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')

我的问题是如何在不将输入设为 tf.placeholder 的情况下将其导入 TensorRT？我能找到的所有示例都使用 tf.placeholder 作为输入。 This example 建议我可以使用 SavedModel class 将迭代器替换为占位符，但我似乎找不到有关如何实现该操作的任何文档。

谢谢！

编辑：感谢下面的帮助，这是我的解决方案

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')

这将写出一个 TensorRT 可以使用的 UFF 文件。我遇到的最大问题是：

我没有意识到 optimize_for_inference_lib.optimize_for_inference 操作将 iterator 替换为 tf.placeholder
我不知道将数据馈送到哪个节点进行评估：您可以将数据馈送到 IteratorGetNext 节点

Answer 1

由于您已经在检查点中保存了经过训练的图，理论上对您来说最简单的解决方案是通过 optimize_for_inference.[=22 导出推理图=]

此工具既适用于已经冻结的图表，也适用于您的情况，即变量仍被定义的图表。假设您采用冻结图方式，第一步是通过以下方式将图形变量转换为常量：

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

这将生成一个名为 frozen_model.pb 的新二进制文件，其中 Variable 操作替换为具有从检查点文件加载的值的 Const 操作。

然后，您需要生成推理图：

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

这会将 IteratorGetNext 节点替换为浮动占位符。您可能想要选择另一个节点，在这种情况下只需更改名称即可。您还可以通过 --placeholder_type_enum 选项更改生成的占位符的类型。在这种情况下，您需要提供一个与 DataType enum.

中所需数据类型相匹配的整数值

注意： 我说 "in theory" 是因为实际上从我做的测试中检查生成的初始图，似乎那里仍然有一些奇怪的操作，但实际上并不是推理所必需的。您可能需要通过 nvidia 的 Graph Surgeon or TF's graph transform tool

进一步处理您的图表

训练后用占位符交换 TensorFlow 数据集输入管道

Swap a TensorFlow Dataset input pipeline with a placeholder after training

python

tensorflow

tensorrt