Question

我有一个自定义估算器。我正在尝试优化此估算器以进行推理（冻结+优化）。

这是我尝试在model_fn()中添加的内容：

input_tensors = [input_1, input_2, input_3]
output_tensors = [output_1, output_2, output_3, output_4]

tmp_g = tf.get_default_graph().as_graph_def()

from tensorflow.python.tools.optimize_for_inference_lib import optimize_for_inference

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    tmp_g = tf.graph_util.convert_variables_to_constants(sess, tmp_g, [n.name[:-2] for n in output_tensors])
    dtypes = [n.dtype for n in input_tensors]

    tmp_g = optimize_for_inference(
                tmp_g,
                [n.name[:-2] for n in input_tensors],
                [n.name[:-2] for n in output_tensors],
                [dtype.as_datatype_enum for dtype in dtypes],
                False)


    output_1, output_2, output_3, output_4 = tf.import_graph_def(tmp_g, input_map={n.name[:-2] + ':0': n for n in input_tensors}, return_elements=['output_1', 'output_2', 'output_3', 'output_4 '])
    from tensorflow.python.estimator.model_fn import EstimatorSpec

    return EstimatorSpec(mode=mode, predictions={
                "output_1": output_1[0],
                "output_2": output_2[0],
                "output_3": output_3[0],
                "output_4": output_4[0]
        })

但是，当调用import_graph_def时，出现以下错误：

ValueError：节点'bert / embeddings / ExpandDims'：连接无效具有1个输出的源节点IteratorGetNext的输出1

（我的输入连接到BERT模型）。

任何想法来自何处？以及如何解决这个问题？

Answer 1

当我致电tf.import_graph_def时，我的张量名称错误。

我用this question寻求帮助。

Here是Colab上的一个最小示例。

我通过执行以下操作解决了该问题：

output_1, output_2, output_3, output_4 = tf.import_graph_def(tmp_g, input_map={n.name: n for n in input_tensors}, return_elements=[n.name for n in output_tensors])

多个输入/输出自定义估计器的推理优化

1 个答案: