ONNX系列四 --- 使用ONNX使TensorFlow模型可移植

目录

TensorFlow简介

安装和导入转换器

快速浏览模型

将TensorFlow模型转换为ONNX

摘要和后续步骤

参考文献

---

• 下载源547.1 KB

系列文章列表如下：

ONNX系列一 --- 带有ONNX的便携式神经网络

ONNX系列二 --- 使用ONNX使Keras模型可移植

ONNX系列三 --- 使用ONNX使PyTorch AI模型可移植

ONNX系列四 --- 使用ONNX使TensorFlow模型可移植

ONNX系列五 --- 在C＃中使用可移植的ONNX AI模型

ONNX系列六 --- 在Java中使用可移植的ONNX AI模型

ONNX系列七 --- 在Python中使用可移植的ONNX AI模型

在关于2020年使用便携式神经网络的系列文章中，您将学习如何将TensorFlow模型转换为便携式ONNX格式。

由于ONNX并非用于构建和训练模型的框架，因此我将首先简要介绍TensorFlow 1.0和TensorFlow 2.0。这对于从头开始并考虑将TensorFlow作为构建和训练其模型的框架的工程师而言非常有用。

TensorFlow简介

TensorFlow最初是由Google Brain团队创建的。1.0版于2017年2月发布。它的体系结构使其可以轻松部署在任何计算设备上。例如，可以将其部署到具有CPU，GPU或TPU的服务器群集。它也可以部署到移动设备和边缘设备。它的灵活部署使其成为生产环境的最爱。但是，由于PyTorch引入了“Define-by-Run”方案，该方案在研究环境中已失去了PyTorch的市场份额，该方案是一种可以根据网络中发生的计算动态定义的神经网络。TensorFlow 1.0基于“Define-by-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使用“检查点”文件来持久化模型。下面的示例指定一个检查点文件。约定是使用扩展名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需要此功能，因为它是一个推理图并且没有变量。tf2onnx工具包含一个名为process_tf_graph的函数，如果您想在代码内进行转换，可以尝试使用该函数。但是，如果最终收到错误消息KeyError: tf.float32_ref，则最好从命令行转换文件。

以下是将TensorFlow 1.0检查点文件转换为ONNX的命令。请注意，您需要找到元文件并将其传递给tf2onnx。您还需要指定输入名称和输出名称。

python -m tf2onnx.convert --checkpoint ./tensorflow/tensorflow_model.ckpt.meta --output tfmodel.onnx --inputs input:0 --outputs output:0

以下是转换TensorFlow 2.0模型的命令。您需要指定目录是用来保存模型到磁盘。（它不会保存在单个文件中。）您还需要指定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低的API的人使用。

接下来，我们将研究如何使用C＃中的ONNX Runtine。

参考文献

• https://www.tensorflow.org/
• https://github.com/onnx/tensorflow-onnx
• https://github.com/keithpij/onnx-lab