KeyError:将Tensorflow模型冻结为UFF图

时间:2018-07-14 15:35:21

标签: python tensorflow protocol-buffers tensorrt

我已经使用Tensorflow Estimator API训练了自定义CNN模型。我已经成功冻结了该图,但是转换为UFF失败并引发以下错误:

frozen_graph_filename = "Frozen_model.pb"
TMP_UFF_FILENAME = "output.uff"
output_name = "sigmoid"

uff_model = uff.from_tensorflow_frozen_model(
    frozen_file=frozen_graph_filename,
    output_nodes=[output_name],
    output_filename=TMP_UFF_FILENAME,
    text=False,
)

执行上述转换的代码:

prefix/OneShotIterator
prefix/IteratorGetNext
prefix/Reshape/shape
prefix/Reshape
prefix/Reshape_1/shape
prefix/Reshape_1
prefix/conv1/kernel
prefix/conv1/bias
.
.
.
prefix/logits/MatMul
prefix/logits/BiasAdd
prefix/sigmoid

图中节点的名称为

tf.graph_util.remove_training_nodes

有没有办法删除前两个Iterator节点?在培训环境之外,它们是无用的。我也使用过// templates/base2.html.twig <!DOCTYPE html> <html> <body> <h2>test</h2> {% block blockoftext %}{% endblock %} </body> </html> ,但并不能缓解我所面临的问题。

1 个答案:

答案 0 :(得分:0)

可以使用Graph Transform Tool来完成我想实现的目的,为了能够使用该工具,必须克隆tensorflow存储库并运行configure文件以设置工作区。请按照1中提供的说明来构建工具。 完成调用后,

bazel-bin/tensorflow/tools/graph_transforms/transform_graph \
--in_graph=tensorflow_inception_graph.pb \
--out_graph=optimized_inception_graph.pb \
--inputs='Mul:0' \
--outputs='softmax:0' \
--transforms='
strip_unused_nodes(type=float, shape="1,299,299,3")
remove_nodes(op=Identity, op=CheckNumerics)
fold_old_batch_norms
'

具有优化的图形后,将其传递给

uff.from_tensorflow_frozen()

https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tools/graph_transforms/README.md 1 https://www.tensorflow.org/mobile/prepare_models#how_do_you_get_a_model_you_can_use_on_mobile