使 TensorFlow 模型可移植到 ONNX
5.00/5 (1投票)
在本文中，我简要概述了 TensorFlow 1.0 和 TensorFlow 2.0，供那些寻找用于构建和训练神经网络的深度学习框架的人参考。
在这一系列关于在 2020 年使用便携式神经网络的文章中，您将学习如何将 TensorFlow 模型转换为便携式 ONNX 格式。
由于 ONNX 不是一个用于构建和训练模型的框架，我将从简要介绍 TensorFlow 1.0 和 TensorFlow 2.0 开始。这对于从零开始并考虑使用 TensorFlow 作为构建和训练模型的框架的工程师来说将非常有用。
TensorFlow 简介
TensorFlow 最初由 Google Brain 团队创建。1.0 版本于 2017 年 2 月发布。它具有易于在任何计算设备上部署的架构。例如，它可以部署到具有 CPU、GPU 或 TPU 的服务器集群。它还可以部署到移动设备和边缘设备。其灵活的部署使其成为生产环境的首选。然而，在研究环境中，它已经失去了 PyTorch 的市场份额，因为 PyTorch 引入了“动态执行”（Define-by-Run）方案，这是一种可以根据网络内发生的计算动态定义的神经网络类型。TensorFlow 1.0 基于“定义后执行”（Define-and-Run）方案，其中网络被定义和固定。运行时唯一发生的事情是将数据输入网络。
为了吸引研究人员，TensorFlow 2.0 引入了“动态执行”（Define-by-Run）方案。此外，TensorFlow 2.0 集成了 Keras。Keras 是一个用于构建神经网络的高级 API，易于使用。TensorFlow 2.0 文档中的大多数 TensorFlow 2.0 示例都使用 Keras，但也可以仅使用 TensorFlow 2.0 构建神经网络。如果您需要对网络进行底层控制，或者您正在将现有的 TensorFlow 1.0 网络迁移到 TensorFlow 2.0，您可能希望这样做。
安装和导入转换器
在将 TensorFlow 模型转换为 ONNX 之前，您需要安装 `tf2onnx` 包，因为它不包含在任何版本的 TensorFlow 中。以下命令将安装此实用程序
pip install tf2onnx
安装后，可以使用下面的导入语句将转换器导入到您的模块中。然而，正如我们在接下来的几节中将看到的，从命令行使用 `tf2onnx` 实用程序更简单。
import tf2onnx
快速浏览模型
TensorFlow 1.0 和 TensorFlow 2.0 都提供了一个低级 API。例如，它们允许您设置权重和偏差以更精确地控制模型训练。但是，为了将 TensorFlow 1.0 和 TensorFlow 2.0 模型转换为 ONNX，我们只关心指定输入、输出以及模型以 TensorFlow 格式保存位置的代码。（本文有一个完整的演示，用于预测 MNIST 数据集中手写样本的数字。）
Tensorflow 1.0 使用占位符方法创建用于指示模型输入和输出的特殊变量。示例如下。为了更轻松地转换为 ONNX，最好在设置占位符时指定名称。此处输入名称为“input”，输出名称为“output”。
# tf Graph input
X = tf.placeholder("float", [batch_size, num_input], name="input")
Y = tf.placeholder("float", [batch_size, num_classes], name="output")
一旦在 TensorFlow 1.0 中创建了带有模型的会话，就可以使用下面的代码保存它。有点令人困惑的是，文件格式在 TensorFlow 1.0 和 TensorFlow 2.0 之间发生了变化。TensorFlow 1.0 使用“checkpoint”文件来持久化模型。下面的示例指定了一个 checkpoint 文件。约定是使用 `.ckpt` 扩展名。将有其他文件保存到与 `.ckpt` 文件相同的目录中，因此最好为保存的模型创建一个单独的目录。
saver = tf.train.Saver()
save_path = saver.save(sess, './tensorflow/tensorflow_model.ckpt')
TensorFlow 2.0 使用“SavedModel”格式，这使得转换过程稍微容易一些。以下是保存 TensorFlow 2.0 模型的命令。请注意，未指定文件名而是指定了目录。（顺便说一句，集成到 TensorFlow 2.0 中的 Keras 使用相同的格式。）
tf.saved_model.save(model, './tensorflow')
将 TensorFlow 模型转换为 ONNX
将 TensorFlow 模型转换为 ONNX 的最简单方法是从命令行使用 `tf2onnx` 工具。当从命令行使用时，`tf2onnx` 会将保存的 TensorFlow 模型转换为另一个表示 ONNX 格式的模型的文件。可以在代码中运行转换，但对于内存中的 TensorFlow 模型，`tf2onnx` 可能会在冻结图时遇到问题。冻结是一个将图中的所有变量转换为常量的过程。这对于 ONNX 是必需的，因为 ONNX 是一个推理图，没有变量。`tf2onnx` 工具包含一个名为 `process_tf_graph` 的函数，如果您想在代码中进行转换，可以尝试使用它。但是，如果您最终收到 `KeyError: tf.float32_ref` 错误消息，那么最好从命令行转换文件。
下面的命令是将 TensorFlow 1.0 checkpoint 文件转换为 ONNX。请注意，您需要找到 `.meta` 文件并将其传递给 `tf2onnx`。您还需要指定输入名称和输出名称。
python -m tf2onnx.convert --checkpoint ./tensorflow/tensorflow_model.ckpt.meta --output tfmodel.onnx --inputs input:0 --outputs output:0
下面的命令是将 TensorFlow 2.0 模型转换为 ONNX。您需要指定保存模型到磁盘的*目录*。（它不是保存在单个文件中。）您还需要指定 ONNX 输出文件。您不需要指定输入名称和输出名称。
python -m tf2onnx.convert --saved-model ./tensorflow --output tfmodel.onnx
总结和后续步骤
在本文中，我简要概述了 TensorFlow 1.0 和 TensorFlow 2.0，供那些寻找用于构建和训练神经网络的深度学习框架的人参考。然后，我展示了如何安装 `tf2onnx` 转换包并将 TensorFlow 模型转换为 ONNX 格式。我还展示了一个很容易出错的地方：指定输入参数的形状。
由于本文的目的是演示将 TensorFlow 模型转换为 ONNX 格式，因此我没有详细介绍构建和训练 TensorFlow 模型。本文的代码示例包含探索 TensorFlow 本身的代码。有一个 TensorFlow 1.0 的演示，拥有现有 TensorFlow 模型的工程师会发现它很有用。还有一个 TensorFlow 2.0 的演示，适用于那些想要比 Keras 更低级 API 的人。
接下来，我们将从 C# 使用 ONNX Runtime。