【深度学习实践】移动端和嵌入式设备上TensorFlow Lite模型部署

将从以下几个部分为您详细介绍 TensorFlow Lite，包括它的基本概念、主要特性、工作流程，以及如何在 Python 和 C++ 中进行推理调用。为了让您更好地理解和实践，我们还会提供示例代码并在重要部分添加注释说明。

---

一、TensorFlow Lite 简介

1.1 什么是 TensorFlow Lite

TensorFlow Lite（以下简称 TFLite）是谷歌为移动端和嵌入式设备推出的一个轻量级的深度学习推理框架。它针对资源受限的环境进行了优化，使得在移动设备（Android、iOS）、物联网设备、微控制器等硬件上部署机器学习模型成为可能。

1.2 TensorFlow Lite 的主要特性

1. 轻量化：TFLite 的目标是尽可能减少对存储空间、计算资源的占用，更好地适配移动和嵌入式设备。
2. 跨平台：TFLite 提供了 C++ API、Java API、Swift API 等，可以在 Android、iOS、Linux 以及其他嵌入式系统中进行推理。
3. 推理速度快：TFLite 提供了多种硬件加速选项（如 GPU、DSP、NN API 等），在移动设备上可以获得更快的推理速度。
4. 易于使用：TensorFlow 原生支持将训练好的模型转换为 TFLite 格式，用户也可使用转换工具对部分其他格式的模型进行转换并部署。

1.3 TensorFlow Lite 的工作流程

大致可以分为以下几个阶段：

1. 训练模型（可选）：
   • 在 TensorFlow（或其他工具）中完成模型训练，得到一个 .pb 或者 SavedModel 格式的模型文件。
2. 模型转换：
   • 使用 tflite_convert 或者新版 TensorFlow Lite Converter（Python API: tf.lite.TFLiteConverter）将模型转换为 .tflite 格式。
   • 在转换的过程中，可以选择是否进行量化（Quantization），例如 8bit 量化、混合量化等，以进一步减小模型大小并加速推理。
3. 加载模型并推理：
   • 在移动端或嵌入式设备，使用 TFLite 的 API（C++、Java、Swift、Python等）加载 .tflite 模型，并进行推理。

---

二、Python 调用 TensorFlow Lite 进行推理的示例

在 Python 环境下，常见的用例是对模型进行快速测试或开发环境模拟，或者在服务器端使用 Python 来调用 TFLite 进行推理。TensorFlow 官方在 Python 下提供了 tensorflow.lite.Interpreter 来进行推理。

下面示例演示了如何：

1. 下载或使用已有的 .tflite 模型（例如，一个简单的图像分类模型 MobileNetV1）。
2. 使用 TFLite 的 Python API 加载并推理。
3. 对输入和输出进行预处理和后处理。

2.1 先准备模型

此处为了演示，您可以在 TensorFlow Hub 下载一个 TFLite 的预训练模型，例如 MobileNetV2 classification 模型（后缀是 .tflite）。假设我们将它命名为 mobilenet_v2_1.0_224.tflite。

2.2 Python 示例代码

环境要求：

• Python 3.x
• tensorflow (确保包含 TFLite 的 Python 依赖，一般安装 tensorflow 就会包含 TFLite Converter 和 Interpreter)

以下示例演示用 tensorflow.lite.Interpreter 进行推理的关键步骤。

import tensorflow as tf
import numpy as np
from PIL import Image

# ========== 第一步：加载 TFLite 模型 ==========
# TFLiteInterpreter 可以通过指定模型文件来初始化
interpreter = tf.lite.Interpreter(model_path="mobilenet_v2_1.0_224.tflite")

# 分配张量空间
interpreter.allocate_tensors()

# ========== 第二步：获取输入和输出张量的信息 ==========
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 输出一下查看（可选）
print("Input details:", input_details)
print("Output details:", output_details)

# 这里假设 input_details[0] 给出的 shape 为 [1, 224, 224, 3]
# 说明模型需要输入一张 224x224，RGB 三通道的图像

# ========== 第三步：准备测试图像 ========== 
# 加载一张图片并 resize 到 224x224
image = Image.open("test_image.jpg").resize((224, 224))
# 转成 numpy array
input_data = np.array(image, dtype=np.float32)

# 由于一般移动端模型会假定输入是 [0,1] 范围或特定均值方差，因此需要做归一化
# 不同模型预处理可能不同，通常可以先将像素值缩放到 [0, 1]
input_data = input_data / 255.0

# 模型输入 shape 为 [1, 224, 224, 3]，所以在最前面加个 batch 维度
input_data = np.expand_dims(input_data, axis=0)

# 将数据设置到 TFLite 的输入张量中
interpreter.set_tensor(input_details[0]['index'], input_data)

# ========== 第四步：执行推理 ==========
interpreter.invoke()

# ========== 第五步：获取推理结果并进行后处理 ==========
output_data = interpreter.get_tensor(output_details[0]['index'])
# output_data.shape -> (1, 1001) 这与模型的类别数相关
# 取出 top1
pred_label = np.argmax(output_data[0])
confidence = output_data[0][pred_label]

print("Predicted label: {}, Confidence: {:.3f}".format(pred_label, confidence))

2.3 代码讲解

1. 加载 Interpreter：

   interpreter = tf.lite.Interpreter(model_path="mobilenet_v2_1.0_224.tflite")
   

   这一步会读取 .tflite 文件，并将模型结构加载到内存中。

2. 分配张量空间：

   interpreter.allocate_tensors()
   

   在 TFLite 推理前，需要分配运行中所需的张量内存空间。

3. 输入输出信息：

   • input_details 和 output_details 包含了张量的具体信息，包括张量的索引（index）、数据类型（dtype）、形状（shape）等。
   • 可以通过 interpreter.get_input_details() 获取模型输入张量的元信息。
   • 可以通过 interpreter.get_output_details() 获取模型输出张量的元信息。

4. 数据预处理：
   由于大多数模型都假定图像数据在 [0,1] 之间或者减去一定均值、除以一定标准差，这里示例采用最简单的将像素值归一化到 [0,1] 的方法。具体预处理流程需参考模型说明。

5. 推理：
   调用 interpreter.invoke() 即可执行一次前向推理。

6. 输出处理：

   • 调用 interpreter.get_tensor(output_details[0]['index']) 可以获取推理结果的张量。
   • 不同模型输出形式不一样，如本例子输出张量的 shape 为 [1, 1001]，表示 1001 个分类结果（包含背景类）。

---

三、C++ 调用 TensorFlow Lite 进行推理的示例

在 C++ 环境（特别是嵌入式环境）下，使用 TensorFlow Lite 的主要流程与 Python 类似：加载模型、准备输入、执行推理、获取输出。但在 C++ 中，需要手动编译和链接 TFLite 库，并管理好内存和数据指针。

下面将介绍一个较为精简的 C++ 示例，用来演示如何加载 .tflite 文件并进行推理。假设我们仍然使用同一个 mobilenet_v2_1.0_224.tflite 模型。

3.1 准备环境

1. 下载 TensorFlow Lite 源码：可以从 tensorflow 源码仓库 中获取。
2. 编译 TFLite 库：根据官方文档，您可以使用 Bazel 或 CMake 编译生成 libtensorflow-lite.a 或者动态库 libtensorflow-lite.so。
3. 包含头文件：需要包含 TFLite C++ API 对应的头文件，如 interpreter.h、model.h、c_api.h 等。
4. 准备第三方库：可能还需要链接其他依赖库，如 pthread、dl 等。

3.2 C++ 示例代码

下面的示例做了尽量详细的注释说明，演示了使用 C++ API 进行推理的核心流程。请根据实际编译环境来进行相应的调整。

#include <iostream>
#include <fstream>
#include <vector>
#include <memory>

// TensorFlow Lite 相关头文件
#include "tensorflow/lite/interpreter.h"
#include "tensorflow/lite/kernels/register.h"
#include "tensorflow/lite/model.h"
#include "tensorflow/lite/c/c_api_types.h"

// 假设我们有一个简单的函数可以读取图片数据并返回 float 向量，
// 实际项目中可使用 OpenCV / stb_image / 自定义读取等方式。
bool LoadImageToBuffer(const std::string& image_path, std::vector<float>& out_img_data, int& width, int& height, int& channels) {
    // 伪代码：此处仅作示例
    // 实际需要用OpenCV 或者其他库将 image 读取并 resize 到 (224,224)
    // 并将其转换为 float 格式归一化到 [0,1] 范围
    // 这里假设 width=224, height=224, channels=3
    width = 224;
    height = 224;
    channels = 3;
    out_img_data.resize(width * height * channels);
    // ... 填充 out_img_data 的像素数据 ...
    return true;
}

int main(int argc, char* argv[]) {
    // ========== 第一步：加载 TFLite 模型 ==========
    // 参数检查
    if (argc < 3) {
        std::cerr << "Usage: " << argv[0] << " <tflite_model> <input_image>\n";
        return -1;
    }
    std::string model_path = argv[1];
    std::string image_path = argv[2];

    // 读取模型文件到 Model
    std::unique_ptr<tflite::FlatBufferModel> model = tflite::FlatBufferModel::BuildFromFile(model_path.c_str());
    if (!model) {
        std::cerr << "Failed to mmap model " << model_path << "\n";
        return -1;
    }

    // ========== 第二步：创建解释器 Interpreter 并分配张量 ==========
    tflite::ops::builtin::BuiltinOpResolver resolver;
    std::unique_ptr<tflite::Interpreter> interpreter;

    tflite::InterpreterBuilder(*model, resolver)(&interpreter);
    if (!interpreter) {
        std::cerr << "Failed to construct interpreter\n";
        return -1;
    }

    // 分配张量
    if (interpreter->AllocateTensors() != kTfLiteOk) {
        std::cerr << "Failed to allocate tensors!\n";
        return -1;
    }

    // ========== 第三步：获取输入张量信息，准备输入数据 ==========
    int input_idx = interpreter->inputs()[0];
    TfLiteTensor* input_tensor = interpreter->tensor(input_idx);
    if (!input_tensor) {
        std::cerr << "Failed to get input tensor\n";
        return -1;
    }

    // 读取图片到缓冲区
    int width, height, channels;
    std::vector<float> img_data;
    if (!LoadImageToBuffer(image_path, img_data, width, height, channels)) {
        std::cerr << "Failed to load image data\n";
        return -1;
    }

    // 检查输入张量形状是否与模型期望相符 (例如 [1,224,224,3])
    // 这里仅简单检查一下
    if (input_tensor->dims->data[1] != height ||
        input_tensor->dims->data[2] != width ||
        input_tensor->dims->data[3] != channels) {
        std::cerr << "Model input shape != loaded image shape.\n";
        return -1;
    }

    // 将图像数据复制给 input_tensor 的 buffer
    float* input_buffer = interpreter->typed_tensor<float>(input_idx);
    memcpy(input_buffer, img_data.data(), img_data.size() * sizeof(float));

    // ========== 第四步：执行推理 ==========
    if (interpreter->Invoke() != kTfLiteOk) {
        std::cerr << "Failed to invoke tflite!\n";
        return -1;
    }

    // ========== 第五步：获取输出并处理 ==========
    int output_idx = interpreter->outputs()[0];
    TfLiteTensor* output_tensor = interpreter->tensor(output_idx);
    if (!output_tensor) {
        std::cerr << "Failed to get output tensor\n";
        return -1;
    }

    // 假设输出是一个 [1,1001] 的分类结果
    float* output_buffer = interpreter->typed_tensor<float>(output_idx);

    // 找到最大置信度的索引
    int pred_label = -1;
    float max_value = -1.0f;
    int num_classes = output_tensor->dims->data[1];
    for (int i = 0; i < num_classes; ++i) {
        if (output_buffer[i] > max_value) {
            max_value = output_buffer[i];
            pred_label = i;
        }
    }

    std::cout << "Predicted label: " << pred_label
              << ", Confidence: " << max_value << std::endl;

    return 0;
}

3.3 代码讲解

1. 构建模型 FlatBufferModel

   std::unique_ptr<tflite::FlatBufferModel> model = tflite::FlatBufferModel::BuildFromFile(model_path.c_str());
   

   读取 .tflite 文件并构建一个 FlatBufferModel 对象，该对象内部包含模型结构和权重参数。

2. InterpreterBuilder

   tflite::ops::builtin::BuiltinOpResolver resolver;
   tflite::InterpreterBuilder(*model, resolver)(&interpreter);
   

   • BuiltinOpResolver 负责将模型中的算子映射到内置的算子实现上。
   • InterpreterBuilder 将 model 和 resolver 组合在一起，生成一个 Interpreter 实例。

3. AllocateTensors

   interpreter->AllocateTensors();
   

   分配模型推理需要的内存空间。

4. 向输入张量填充数据

   float* input_buffer = interpreter->typed_tensor<float>(input_idx);
   memcpy(input_buffer, img_data.data(), img_data.size() * sizeof(float));
   

   • 通过 typed_tensor<float>(input_idx) 获取可写的浮点指针。
   • 将预处理后的图像数据拷贝到输入张量的内存中。

5. 执行推理

   interpreter->Invoke();
   

   完成一次前向推理计算。

6. 获取输出

   • 通过 interpreter->outputs()[0] 获取第一个输出张量的索引，然后获取对应的 TfLiteTensor*。
   • 再使用 typed_tensor<float>(output_idx) 拿到输出浮点数组。

7. 后处理

   • 遍历输出数组，找到最大值的索引（分类任务中通常这样确定预测类别）。
   • 打印结果。

3.4 编译方式示例

编译上述代码时，需要链接 tensorflow-lite 静态/动态库以及其他依赖库。以一个简化的 g++ 命令为例（假设已经生成了 libtensorflow-lite.a 并放在某个路径）：

g++ -std=c++11 -I/path/to/tensorflow/lite/headers \
    -I/path/to/absl/headers \
    -L/path/to/tflite/library \
    my_tflite_inference.cpp \
    -ltensorflow-lite -labsl_base -labsl_malloc_internal -lpthread \
    -o tflite_inference

实际需要根据您的具体环境调整头文件路径和库文件路径。

---

四、从 TensorFlow 模型到 TFLite 模型的转换

在 Python 环境下，如您拥有一个 TensorFlow 训练好的 SavedModel 或者 Keras 模型，可以通过以下方式进行转换：

import tensorflow as tf

# 例如：加载一个 Keras 模型
keras_model = tf.keras.applications.MobileNetV2(weights='imagenet')

# 将 Keras 模型转换成 TFLite 模型
converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
# 可选：指定量化策略
# converter.optimizations = [tf.lite.Optimize.DEFAULT]

tflite_model = converter.convert()

# 保存到文件
with open("mobilenet_v2_1.0_224.tflite", "wb") as f:
    f.write(tflite_model)

转换完成后，就可以拿到 .tflite 文件并在移动端或嵌入式设备使用。

---

五、总结

1. TensorFlow Lite 适用于在移动端或嵌入式设备部署深度学习模型，具有轻量、高效的特点。
2. 在 Python 中使用非常简便：
   • 通过 tf.lite.Interpreter 加载 .tflite 文件。
   • 设置输入数据并调用 invoke() 进行推理。
   • 获取输出张量进行后处理。
3. 在 C++ 中，需要将 TensorFlow Lite 编译为静态或动态库，然后：
   • 加载 .tflite 文件并构建 FlatBufferModel。
   • 创建 Interpreter 并分配张量。
   • 填充输入数据并调用 Invoke() 推理。
   • 获取输出并进行相应的后处理。
4. 如果您需要从 TensorFlow 或 Keras 模型转换到 .tflite，可以使用 TFLite Converter（tflite_convert 命令或 Python API）完成。

通过上面的示例和说明，您已经了解了 TensorFlow Lite 的工作流程以及在 Python 与 C++ 下如何调用推理。接下来，您可以根据自身项目需求，对预处理、后处理、模型量化、硬件加速等做进一步的优化和扩展。祝您在实践中一切顺利！