YOLOv8 在单片机上的几种部署方案

news2025/5/20 21:05:34

在这里插入图片描述

YOLOv8 在单片机上的部署方案

单片机资源（如内存、计算能力）有限，直接部署完整的 YOLOv8 模型并不现实。不过，我们可以通过模型量化、优化和使用轻量级框架来实现简化版的目标检测。下面为你介绍几种可行的方案：

方案一：使用 TensorFlow Lite Micro + YOLOv8 简化模型

1. 模型转换与优化

首先在 PC 上对 YOLOv8 进行简化和量化：

import torch
from ultralytics import YOLO
import tensorflow as tf
from onnx_tf.backend import prepare

# 加载 YOLOv8 模型
model = YOLO("yolov8n.pt")  # 使用 Nano 版本

# 导出为 ONNX 格式
model.export(format="onnx", imgsz=(320, 320))  # 减小输入尺寸

# 转换 ONNX 到 TensorFlow
import onnx
onnx_model = onnx.load("yolov8n.onnx")
tf_rep = prepare(onnx_model)
tf_rep.export_graph("yolov8n_tf")

# 转换为 TensorFlow Lite 并应用量化
converter = tf.lite.TFLiteConverter.from_saved_model("yolov8n_tf")
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_quant_model = converter.convert()

# 保存 TFLite 模型
with open("yolov8n_quant.tflite", "wb") as f:
    f.write(tflite_quant_model)

2. 在单片机上部署 TensorFlow Lite Micro

以 Arduino Nano 33 BLE Sense 为例：

#include "tensorflow/lite/micro/all_ops_resolver.h"
#include "tensorflow/lite/micro/micro_error_reporter.h"
#include "tensorflow/lite/micro/micro_interpreter.h"
#include "tensorflow/lite/schema/schema_generated.h"
#include "model_data.h"  // 包含量化后的 YOLOv8 模型

// 定义输入输出张量
const int kInputTensorIndex = 0;
const int kOutputTensorIndex = 0;

// 初始化错误报告器
tflite::MicroErrorReporter micro_error_reporter;
const tflite::ErrorReporter* error_reporter = &micro_error_reporter;

// 初始化算子解析器
tflite::AllOpsResolver resolver;

// 加载模型
const tflite::FlatBufferModel* model = 
    tflite::FlatBufferModel::BuildFromBuffer(model_data, model_data_len);

// 创建解释器
constexpr int tensor_arena_size = 136 * 1024;
uint8_t tensor_arena[tensor_arena_size];
tflite::SimpleTensorAllocator tensor_allocator(tensor_arena, tensor_arena_size);
tflite::MicroInterpreter interpreter(model, resolver, &tensor_allocator, error_reporter);

// 分配张量
TfLiteStatus allocate_status = interpreter.AllocateTensors();
if (allocate_status != kTfLiteOk) {
  Serial.println("Failed to allocate tensors!");
  return;
}

// 获取输入输出张量
TfLiteTensor* input_tensor = interpreter.input(kInputTensorIndex);
TfLiteTensor* output_tensor = interpreter.output(kOutputTensorIndex);

// 图像预处理函数（示例）
void preprocess_image(uint8_t* image_data, float* input_data) {
  // 调整图像大小为模型输入尺寸 (320x320)
  // 归一化像素值到 [0, 1] 或 [-1, 1]
  // ...
}

// 后处理函数（简化版 NMS）
void postprocess(float* output_data, int width, int height) {
  // 解析模型输出，提取边界框、类别和置信度
  // 应用非极大值抑制(NMS)
  // ...
}

void setup() {
  Serial.begin(115200);
  // 初始化摄像头
  // ...
}

void loop() {
  // 捕获图像
  uint8_t* image_data = capture_image();
  
  // 预处理图像
  preprocess_image(image_data, input_tensor->data.f);
  
  // 运行推理
  TfLiteStatus invoke_status = interpreter.Invoke();
  if (invoke_status != kTfLiteOk) {
    Serial.println("Failed to invoke interpreter!");
    return;
  }
  
  // 后处理结果
  postprocess(output_tensor->data.f, 320, 320);
  
  // 显示或发送结果
  // ...
  
  delay(100);
}

方案二：使用 TinyML 框架（如 NCNN）

NCNN 是专为移动设备优化的轻量级神经网络推理框架，非常适合单片机：

1. 模型转换

将 YOLOv8 转换为 NCNN 格式：

# 首先将 YOLOv8 导出为 ONNX
yolo export model=yolov8n.pt format=onnx imgsz=320

# 使用 onnx2ncnn 工具转换为 NCNN 格式
onnx2ncnn yolov8n.onnx yolov8n.param yolov8n.bin

# 优化模型
ncnnoptimize yolov8n.param yolov8n.bin yolov8n-opt.param yolov8n-opt.bin 1

2. 在单片机上集成 NCNN

以下是一个简化的 NCNN 集成示例：

#include "net.h"
#include "benchmark.h"
#include "mat.h"

// 初始化网络
ncnn::Net yolov8;
yolov8.load_param("yolov8n-opt.param");
yolov8.load_model("yolov8n-opt.bin");

// 目标检测函数
std::vector<Object> detect_yolov8(const cv::Mat& bgr, float prob_threshold = 0.25f, float nms_threshold = 0.45f)
{
    int img_w = bgr.cols;
    int img_h = bgr.rows;

    // 图像预处理
    ncnn::Mat in = ncnn::Mat::from_pixels_resize(bgr.data, ncnn::Mat::PIXEL_BGR, bgr.cols, bgr.rows, 320, 320);
    
    // 归一化
    const float mean_vals[3] = {0.f, 0.f, 0.f};
    const float norm_vals[3] = {1/255.f, 1/255.f, 1/255.f};
    in.substract_mean_normalize(mean_vals, norm_vals);

    // 运行推理
    ncnn::Extractor ex = yolov8.create_extractor();
    ex.set_num_threads(2);
    ex.input("images", in);

    ncnn::Mat out;
    ex.extract("output", out);

    // 后处理
    std::vector<Object> objects;
    // ... 解析输出并应用 NMS
    
    return objects;
}

void setup() {
  // 初始化串口和摄像头
}

void loop() {
  // 捕获图像
  cv::Mat image = capture_image();
  
  // 检测目标
  std::vector<Object> objects = detect_yolov8(image);
  
  // 处理检测结果
  // ...
  
  delay(100);
}

方案三：使用 YOLO-NAS Tiny

YOLO-NAS 是一种较新的轻量级目标检测模型，性能优于 YOLOv5/YOLOv8 的 Nano 版本：

# 安装 super-gradients
pip install super-gradients

# 导出 YOLO-NAS Tiny 为 ONNX
from super_gradients.training import models

# 加载模型
model = models.get("yolo_nas_s", pretrained_weights="coco")

# 导出为 ONNX
model.export("yolo_nas_s.onnx", input_shape=(3, 320, 320))

# 然后按照上述方法将 ONNX 转换为适合单片机的格式