在 Android 上使用 TensorFlow Lite 的 ONNX 神经网络模型





5.00/5 (1投票)
在本文中,我们将采用一个预训练的神经网络,并将其适配到TensorFlow Lite中使用。
这是关于在Android上使用TensorFlow Lite系列文章的第二篇。在前一篇文章中,我们为使用TensorFlow模型和Android开发设置了开发环境。在这里,我们将继续适配神经网络以用于TensorFlow Lite。
在移动设备上使用TensorFlow Lite时,您可能希望使用一个已经训练好的网络。预训练网络可以帮助您跳过收集和准备大量数据来训练视觉识别器所需的过程,从而更快地投入生产。预训练网络可能使用多种不同的技术构建。
为了支持神经网络模型保存格式之间的互操作性,存在一种中间格式。这种格式被称为Open Neural Network Exchange格式(ONNX)。ONNX得到了包括Microsoft、Facebook、AMD、IBM、ARM、NVIDIA、Qualcomm等众多公司的支持。
借助ONNX,开发人员可以使用他们偏好的软件和框架来生成他们的神经网络模型,并与可能使用其他AI技术的用户共享。这意味着TensorFlow Lite不仅限于使用仅用TensorFlow实现的模型。可以使用在机器学习框架中共有的运算的网络可以在使用ONNX的ML技术之间共享。请注意,使用实验性或不太常见的运算的网络可能与其他框架的兼容性有限。
在本文中,我将使用一个来自ONNX Model Zoo的模型。在这个仓库中,您会找到用于视觉识别、语音识别、手势分析等的模型。我将下载一个YOLO(You Only Look Once)模型,这是一个视觉分类模型。我使用的是YOLO4实现。在该仓库的`/vision/object_detection_segmentation/yolo4/models`目录下可以找到它。
如果您想使用此仓库中的模型,您将需要使用GIT LFS(Large File System)。使用GIT LFS,您可以克隆一个仓库,然后仅下载感兴趣的文件,而无需下载所有大文件。在克隆此仓库之前,请使用以下命令激活GIT LFS:
git lfs install
GIT会显示LFS已启用的确认信息。接下来,克隆ONNX Model Zoo仓库:
git clone https://github.com/onnx/models
克隆的仓库将包含大型.onnx文件的占位符文件。要下载文件的实际数据,请使用git lfs pull
命令,并传递相对于仓库的文件或文件夹路径以供下载:
git lfs pull --include="vision/object_detection_segmentation/yolov4/model" --exclude=""
或者,如果您有足够的空间、网络带宽和时间,并且希望下载模型中的所有文件,您可以将通配符字符传递给要下载的文件名:
git lfs pull --include="*" --exclude=""
在选择网络时,您还需要查看网络处理的数据的先决条件和后置条件信息。对于某些网络,您可以在模型旁边找到这些信息,通常在一个readme文件中。对于其他一些网络,您可能需要查看其原始来源来查找此信息。视觉识别网络可能需要将图像缩放到特定范围,或者需要以某种方式编码像素值(例如,整数值、浮点值,范围为0.0-1.0、0.0-255.0,或其他范围)。网络输出的数据将是数字数组的形式。请查阅模型的文档,以了解这些数字的含义以及它们的排列方式。对于我选择的网络,在其readme.md文件中有一项说明,图像必须是416x416像素,并且是浮点值。
Netron实用程序虽然不能替代模型文档,但对于从模型中提取信息非常有用。Netron显示了网络中神经元的排列以及其他一些信息。目前,我们感兴趣的部分是网络的输入和输出。使用Netron打开一个YOLO网络模型,然后选择输入节点,会看到如下屏幕:
从这个屏幕我可以看到,这个模型的输入是RGB图像,已经缩放到416x416像素。像素被编码为浮点值。这张图像是NCHW图像。有多种方式可以组织图像数据。
- N - 输入图像批次的数量。通常为1(单个批次)或None,表示批次数量是动态的,可能会变化。
- C - 图像通道。在这种情况下,通道是红色、绿色和蓝色。
- H - 图像的高度。沿着这个维度遍历会从图像的一行到另一行。
- W - 图像的宽度。沿着这个维度遍历会让我们移动到相邻的像素。
Netron的这个屏幕也显示,这个神经网络模型的输出是一个张量,包含为每个输入批次([None x 125 x 13 x 13])的125x13x13个浮点值。这个工具没有指定这些值各自的含义,但我们需要知道它们才能配置输入和输出值。
要将其中一个ONNX模型转换为TensorFlow冻结图,请在终端中使用onnx-tf
,其参数匹配以下模式:
onnx-tf convert -i source_model.onnx -o output_model.pb
片刻之后,您将得到转换后的TensorFlow冻结图。我们实际需要的是一个TensorFlow Lite文件。要将文件从TensorFlow转换为TensorFlow Lite,我使用了一个Python脚本。存在一个命令行工具,但我发现使用Python脚本更加灵活。
import tensorflow as tf
saved_model_dir='/dev/projects/models'
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tf_lite_model = converter.convert()
open('output.tflite', 'wb').write(tf_lite_model)
输出文件的名称可以任意命名,只要它带有.tflite扩展名即可。请确保正确设置扩展名,因为稍后它很重要。
后续步骤
既然我们已经将ONNX模型转换为Tensorflow Lite,就可以在Android应用程序中使用它了。这将在下一篇文章中介绍。