3. 端侧图像分类

任务目标

在Raspberry Pi Pico（RP2040，264KB SRAM）上部署MobileNetV2，实现端侧图像分类。由于MobileNetV2原始模型约13MB，远超Pico存储与内存，本案例采用知识蒸馏后的1.0×MobileNetV2（输入96×96×3，量化至int8），模型大小约1.2MB，运行时峰值内存约180KB。最终在Pico上完成对96×96 RGB图像的实时分类推理。

环境与依赖

  • 硬件：Raspberry Pi Pico（RP2040，2MB Flash）
  • SDK：pico-sdk v1.5+，编译工具链arm-none-eabi-gcc
  • 模型转换：TensorFlow 2.13 + TFLite Converter（量化int8）
  • 推理框架：TFLite Micro（从tensorflow/tflite-micro仓库获取，版本对应TF 2.13）
  • 数据集：ImageNet子集（蒸馏训练用，本案例不涉及训练细节）

模型准备与转换

使用TensorFlow将训练好的MobileNetV2浮点模型转换为int8量化TFLite模型。关键转换参数如下（完整训练代码不在本案例范围）：

import tensorflow as tf

# 加载训练好的浮点模型（输入96x96x3）
converter = tf.lite.TFLiteConverter.from_saved_model('mobilenetv2_96_saved_model')
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_dataset_gen  # 校准数据集生成器
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.int8
converter.inference_output_type = tf.int8
tflite_model = converter.convert()

# 保存模型
with open('mobilenetv2_96_int8.tflite', 'wb') as f:
    f.write(tflite_model)

转换后模型约1.2MB，需通过xxd -i mobilenetv2_96_int8.tflite > model_data.cc生成C数组嵌入固件。

关键完整代码

以下为Pico上TFLite Micro推理核心实现（main.c）：

#include <stdio.h>
#include <pico/stdlib.h>
#include "tensorflow/lite/micro/micro_interpreter.h"
#include "tensorflow/lite/micro/micro_mutable_op_resolver.h"
#include "tensorflow/lite/schema/schema_generated.h"
#include "model_data.cc"  // 包含模型数组

// 输入输出缓冲区（int8）
constexpr int kTensorArenaSize = 190 * 1024;
uint8_t tensor_arena[kTensorArenaSize];

// 图像数据（96x96x3，int8，范围-128~127）
extern const int8_t test_image[96*96*3];  // 由外部提供

int main() {
    stdio_init_all();
    sleep_ms(2000);  // 等待串口就绪

    // 加载模型
    const tflite::Model* model = tflite::GetModel(g_model);
    if (model->version() != TFLITE_SCHEMA_VERSION) {
        printf("Model version mismatch!\n");
        return -1;
    }

    // 注册算子（仅需MobileNetV2用到的）
    static tflite::MicroMutableOpResolver<10> resolver;
    resolver.AddConv2D();
    resolver.AddDepthwiseConv2D();
    resolver.AddAdd();
    resolver.AddMul();
    resolver.AddRelu();
    resolver.AddSoftmax();
    resolver.AddReshape();
    resolver.AddMean();
    resolver.AddPad();
    resolver.AddQuantize();

    // 构建解释器
    static tflite::MicroInterpreter interpreter(
        model, resolver, tensor_arena, kTensorArenaSize);
    if (interpreter.AllocateTensors() != kTfLiteOk) {
        printf("AllocateTensors failed!\n");
        return -1;
    }

    // 获取输入输出张量
    TfLiteTensor* input = interpreter.input(0);
    TfLiteTensor* output = interpreter.output(0);

    // 检查输入尺寸
    if (input->dims->size != 4 || input->dims->data[1] != 96 ||
        input->dims->data[2] != 96 || input->dims->data[3] != 3) {
        printf("Unexpected input shape!\n");
        return -1;
    }

    // 拷贝图像数据到输入张量
    memcpy(input->data.int8, test_image, 96*96*3*sizeof(int8_t));

    // 执行推理
    if (interpreter.Invoke() != kTfLiteOk) {
        printf("Invoke failed!\n");
        return -1;
    }

    // 获取输出（1000类，int8）
    int8_t* out_data = output->data.int8;
    float scale = output->params.scale;
    int zero_point = output->params.zero_point;

    // 反量化并找Top-1
    int max_idx = 0;
    float max_score = -128.0f;
    for (int i = 0; i < 1000; i++) {
        float score = (out_data[i] - zero_point) * scale;
        if (score > max_score) {
            max_score = score;
            max_idx = i;
        }
    }

    printf("Top-1 class: %d, score: %.4f\n", max_idx, max_score);
    return 0;
}

CMake配置

cmake_minimum_required(VERSION 3.13)
include(pico_sdk_import.cmake)
project(mobilenetv2_pico)

# TFLite Micro源码路径（需提前克隆）
set(TFLITE_MICRO_DIR ${CMAKE_CURRENT_SOURCE_DIR}/tflite-micro)

add_executable(mobilenetv2_pico
    main.c
    ${TFLITE_MICRO_DIR}/tensorflow/lite/micro/all_ops_resolver.cc
    # 此处需列出所有编译的.cc源文件（约30+个）
)

target_include_directories(mobilenetv2_pico PRIVATE
    ${TFLITE_MICRO_DIR}
    ${CMAKE_CURRENT_SOURCE_DIR}
)

target_link_libraries(mobilenetv2_pico PRIVATE pico_stdlib)

pico_add_extra_outputs(mobilenetv2_pico)

完整TFLite Micro源文件列表较长，建议使用其官方CMakeLists.txt中的target_sources定义，或直接引用micro_speech示例的构建方式。

编译与烧录

mkdir build && cd build
cmake ..
make -j4
# 生成mobilenetv2_pico.uf2，按住Pico BOOTSEL键后连接USB，拷贝uf2文件到弹出的U盘

验证方法

使用串口监视器（如minicom，波特率115200）观察输出。通过条件：

  • 程序启动后2秒内打印Top-1 class: xxx, score: x.xxxx，无错误信息
  • 使用已知测试图像（如ImageNet中类别为"goldfish"的96×96图像），预期输出类别索引与真实标签一致（goldfish对应索引1）
  • 推理耗时约1.5~2.5秒（可通过time_us_32()测量，本代码未包含计时，如需验证性能可自行添加）

若输出类别与预期不符，检查：模型输入是否需归一化至[-128,127]、图像数据是否按RGB顺序排列、模型是否成功量化。

← 上一章 📖 返回目录 下一章 →