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

Posted

技术标签:

【中文标题】训练后用占位符交换 TensorFlow 数据集输入管道【英文标题】:Swap a TensorFlow Dataset input pipeline with a placeholder after training 【发布时间】:2018-09-13 18:27:40 【问题描述】:

我正在使用新的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 类将迭代器替换为占位符,但我似乎找不到任何有关如何完成此操作的文档。

谢谢!

编辑:感谢以下帮助,这是我的解决方案

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 节点

【问题讨论】:

您是否受限于使用这种相对较低级别的训练方案,或者您是否考虑使用tf.estimator 进行训练和导出的解决方案也可行? 如果 tf.estimator 可以与 TensorRT 一起使用,我很乐意使用它。不过我不确定它会不会。 据我了解,TensorRT 可以处理经过训练的图。 Estimator 可帮助您训练和导出图形,然后将其冻结并将其导入 TensorRT。我明天才能正确回答,我会尝试解释如何设置训练和导出所需的各种代码 您好 GPhilo。我有一个训练有素的图,但它需要一个迭代器。我相信,如果我能理解如何用新数据提供图表,或者在训练后用tf.placeholder 替换迭代器,我就会做好准备。不过,我很想看看如何使用估算器进行设置。 好的,如果你只需要这些,看看这个:github.com/tensorflow/tensorflow/blob/master/tensorflow/python/… 你需要传递应该在你的检查点文件夹中的 graph.txt 文件,确保阅读顶部的cmets关于使用非冻结图 【参考方案1】:

由于您已经在检查点中保存了经过训练的图,理论上最简单的解决方案是通过optimize_for_inference 导出推理图。

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

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 中提供一个与您想要的数据类型相匹配的整数值。

注意:我说“理论上”是因为实际上检查了从测试中生成的初始图,我认为那里仍然有一些奇怪的操作,这些操作对于推理来说并不是真正必要的。您可能需要通过 nvidia 的 Graph Surgeon 或 TF 的 graph transform tool 进一步处理您的图表

【讨论】:

感谢您的回答。我最大的问题是不知道:1)向哪个 TensorFlow 节点提供数据(这里是 IteratorGetNext)和 2)没有意识到 optimize_for_inference 步骤实际上用占位符替换了 IteratorGetNext 是的,对于那里的所有示例和文档,我觉得训练和实际使用网络之间的步骤仍然记录得很差。至少有三件事要做,其中大多数需要一堆不明显的步骤......我希望他们能简化一点:/ 同意。我编辑了我的问题,给出了推理和创建 TensorRT 的 UFF 文件的代码。感谢您的帮助! @GPhilo 我尝试了optimize_for_inference,但它没有用占位符替换IteratorGetNext 节点。我错过了一些标志吗?我的意思是,有一个新的占位符,但数据集加载器节点似乎仍然存在于图中,当我尝试使用编译器前端转换模型时,它在操作 IteratorGetNextIteratorV2 上引发异常 你用的是什么版本的TF? optimize_for_inference 变得有点过时了,它可能不适用于引入的新修改(我看到一个 GitHub 问题,他们建议使用 transform_graph 代替)

以上是关于训练后用占位符交换 TensorFlow 数据集输入管道的主要内容,如果未能解决你的问题,请参考以下文章

TensorFlow 创建 Ai,错误:您必须为占位符张量“input_1/X”提供值

tensorflow 在 GPU 内存上存储训练数据

尝试填充占位符时出现Tensorflow错误

sess.run() 将值输入占位符张量的问题

TensorFlow 占位符依赖于其他占位符

TensorFlow,“‘模块’对象没有属性‘占位符’”