使用TensorFlow Lite将深度学习模型部署到IOT系统

TensorFlow Lite简介

移动设备深度学习框架是部署在手机或者树莓派等小型移动设备上的深度学习框架，可以使用训练好的模型在手机等设备上完成推理任务。这一类框架的出现，可以使得一些推理的任务可以在本地执行，不需要再调用服务器的网络接口，大大减少了预测时间。

TensorFlow Lite是TensorFlow针对移动和嵌入式设备的轻量级解决方案（深度学习框架）。它支持设备内机器学习推理，具有低延迟和小二进制大小。TensorFlow Lite还支持Android神经网络API的硬件加速 。
TensorFlow Lite使用许多技术来实现低延迟，例如优化移动应用程序的内核，预融合激活以及允许更小和更快（定点数学）模型的量化内核。

TensorFlow Lite使用流程

要了解TensorFlow Lite 的使用流程，首先来看一下官方给出的关于该框架的架构图就一目了然了：

tensorflow lite设备端训练 tensorflow lite部署_深度学习

从训练好的TensorFlow模型开始，我们使用TensorFlow Lite转换器将正常训练的模型转换为TensorFlow Lite文件格式。然后，我们就可以在移动应用程序中使用该转换后的文件。

理解TensorFlow Lite模型的在IOT系统中的架构：

  • Java API：围绕Android上的C ++ API的便捷包装器。
  • C ++ API：加载TensorFlow Lite模型文件并调用解释器。Android和iOS都提供相同的库。
  • （InterPreter）解释器：使用一组内核执行模型。解释器支持选择性内核加载; 没有内核它只有100KB，加载了所有内核300KB。这比TensorFlow Mobile要求的1.5M显着降低。
    在部分Android设备上，Interpreter将使用Android神经网络API进行硬件加速，如果没有，则默认为CPU执行。

通过对架构图的理解，现在我们明白了使用TensorFlow Lite 将深度学习模型部署到IOT系统无非分为如下几步：

  • 1、使用TensorFlow Lite转换器将模型转化为.tflite格式；
  • 2、使用相应的应用程序开发工具（这里以Android Studio为例），配置好TensorFlow Lite 模型。
  • 3、将模型和Label放入开发应用程序的项目中
  • 4、界面的开发与API调用
  • 5、将程序打包导出便完成了模型的部署，此时程序便可运行在不同设备相同系统中了。

下面我们根据以上流程一步步分析。

转换模型

下载模型

这里的模型可以是我们自己训练的模型，也可以是别人训练好的模型，这里我使用了TensorFlow官方的预训练模型举例说明：TensorFlow预训练模型

上面提供的模型同时也包括了tflite模型，我们可以直接拿来使用，但是我们也可以使用其他格式的模型来转换。比如我们下载一个MobileNet_v2_1.4_224，解压之后获得以下文件：

mobilenet_v2_1.4_224.ckpt.data-00000-of-00001 
mobilenet_v2_1.4_224.ckpt.meta	mobilenet_v2_1.4_224_frozen.pb
mobilenet_v2_1.4_224.tflite
mobilenet_v2_1.4_224.ckpt.index	
mobilenet_v2_1.4_224_eval.pbtxt	mobilenet_v2_1.4_224_info.txt
mobilenet_v2_1.4_224.tgz

转换模型

网上很多资料是用TensorFlow编译的工具来转化模型，但这里我选择采用有python代码来转换模型，这样配置更简单，使用更方便，配置文件frozen_to_tflite.py如下：

# -*- coding:utf-8 -*-
import tensorflow as tf

in_path = "./mobilenet_v2_1.4_224/mobilenet_v2_1.4_224_frozen.pb"
out_path = "./mobilenet_v2_1.4_224.tflite"
# out_path = "./model/quantize_frozen_graph.tflite"

# 模型输入节点
input_tensor_name = ["input"]
input_tensor_shape = {"input":[1,224,224,3]} # 输入图片的大小
# 模型输出节点
classes_tensor_name = ["MobilenetV2/Predictions/Reshape_1"]

converter = tf.lite.TFLiteConverter.from_frozen_graph(in_path,
                                     input_tensor_name, classes_tensor_name,
                                     input_shapes = input_tensor_shape)
converter.post_training_quantize = True # 定量转换，这样能使得到的模型更小
tflite_model = converter.convert()

with open(out_path, "wb") as f:
    f.write(tflite_model)

执行上面的文件，便生成转换后的模型mobilenet_v2_1.4_224.tflite。
转换前后模型大小的对比：

模型名

转换前

mobilenet_v2_1.4_224_frozen.pb

23.3M

mobilenet_v2_1.4_224.tflite

5.9M

可见经过转换后的模型大小缩小了近4倍。

在项目中配置TensorFlow Lite 框架

  • 创建完成之后，在app目录下的build.gradle配置文件加上以下配置信息：
    在dependencies下加上包的引用，第一个是图片加载框架Glide，第二个就是我们这个项目的核心TensorFlow Lite：
implementation 'com.github.bumptech.glide:glide:4.3.1'
    implementation 'org.tensorflow:tensorflow-lite:0.0.0-nightly'
  • 然后在android下加上以下代码，这个主要是限制不要对tensorflow lite的模型进行压缩，压缩之后就无法加载模型了：
//set no compress models
    aaptOptions {
        noCompress "tflite"
    }

在项目中配置.tflite格式的模型文件

在main目录下创建assets文件夹，这个文件夹主要是存放tflite模型和label名称文件。将.tflite和label文件放入目录内。

加载模型，API的调用，输入图片进行预测

一下是将模型部署在Android手机上，在手机端离线测试的结果：

图一：检测物体：键盘，检测结果：键盘，概率：0.606204，用时：205ms

tensorflow lite设备端训练 tensorflow lite部署_API_02

图二：检测物体：笔记本电脑，检测结果：笔记本电脑，概率：0.410，用时：212ms

tensorflow lite设备端训练 tensorflow lite部署_API_03

图三：检测物体：高脚杯，检测结果：高脚杯，概率：0.215，用时：200ms

tensorflow lite设备端训练 tensorflow lite部署_tensorflow_04

图三：检测物体：鼠标，检测结果：鼠标，概率：0.785，用时：213ms

tensorflow lite设备端训练 tensorflow lite部署_tensorflow_05

从测试结果来看：使用TensorFlow Lite框架部署的mobilenet_v2_1.4_224.tflite模型，具有较高的准确率，平均预测每一张图片需要消耗的时间在200ms左右。