KeyError:冻结的Tensorflow模型到UFF图

Question

我已经用Tensorflow Estimator API训练了一个自定义的CNN模型。我已经成功地冻结了图形，但转换到UFF失败并抛出以下错误：

'KeyError: u'IteratorGetNext:1'

执行上述转换的代码：

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

那么有没有办法删除前两个Iterator节点呢？它们在训练环境之外是无用的。我也用过tf.graph_util.remove_training_nodes，但它并不能缓解我面临的问题。