TensorFlow2.0 中的 XLA - 冻结模型？答案

【问题标题】：XLA in TensorFlow2.0 - frozen model?TensorFlow2.0 中的 XLA - 冻结模型？
【发布时间】：2019-11-29 16:36:20
【问题描述】：

我遵循 XLA AOT 编译的官方指南 (https://www.tensorflow.org/xla/tfcompile)，编译示例工作正常（在 aot/tests 内）。

但是后来我想编译一些稍微大一点的模型，但出现了一个问题：如果 XLA AOT 需要冻结图作为输入（正如我从指南中理解的那样）并且 TensorFlow 2 不再支持冻结图，那么输入会做什么XLA 现在期待吗？

【问题讨论】：

【解决方案1】：

似乎仍有方法可以在 TensorFlow 2 中冻结图表。我按照这篇文章创建了一个冻结图表，然后编译它：https://leimao.github.io/blog/Save-Load-Inference-From-TF2-Frozen-Graph/

# Convert Keras model to ConcreteFunction
full_model = tf.function(lambda x: model(x))
full_model = full_model.get_concrete_function(
    tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype))

# Get frozen ConcreteFunction
frozen_func = convert_variables_to_constants_v2(full_model)
frozen_func.graph.as_graph_def()

layers = [op.name for op in frozen_func.graph.get_operations()]
print("-" * 50)
print("Frozen model layers: ")
for layer in layers:
    print(layer)

print("-" * 50)
print("Frozen model inputs: ")
print(frozen_func.inputs)
print("Frozen model outputs: ")
print(frozen_func.outputs)

# Save frozen graph from frozen ConcreteFunction to hard drive
tf.io.write_graph(graph_or_graph_def=frozen_func.graph,
                  logdir="./frozen_models",
                  name="frozen_graph.pb",
                  as_text=False)

【讨论】：