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

YOLOv8 在单片机上的部署方案

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

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

1. 模型转换与优化

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

python 复制代码
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 为例:

cpp 复制代码
#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 格式:

bash 复制代码
# 首先将 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 集成示例:

cpp 复制代码
#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 版本:

python 复制代码
# 安装 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 转换为适合单片机的格式

资源限制与优化建议

  1. 模型选择:优先选用 Nano 或 Tiny 版本的模型
  2. 输入尺寸:使用较小的输入尺寸(如 160×160 或 320×320)
  3. 量化:使用 8 位或 16 位量化,甚至二值化
  4. 层融合:利用框架的层融合功能减少计算量
  5. 算法简化:只检测关键类别,降低模型复杂度

对于资源极其有限的单片机(如 Arduino Uno),可能需要使用更轻量级的算法,如 Tiny YOLO 或专门为 MCU 设计的目标检测模型。

常见问题及解决办法

CUDA 相关问题:

要保证你的 GPU 驱动版本与 CUDA 版本兼容

可以使用 nvidia-smi 命令查看 GPU 信息

显示问题:

如果你在服务器上运行,可能会遇到无法显示图像的问题,这时可以加上 save=True 参数将结果保存下来。

依赖冲突问题:

可以尝试在全新的虚拟环境中重新安装所有依赖。

如果在部署过程中遇到特定问题,请提供详细的错误信息,以便进一步排查。

优化建议与注意事项

模型压缩策略:

1、使用 YOLOv8 Nano 或定制更小的模型

2、降低输入分辨率(128×128 或 160×160)

3、应用 INT8 或二值化量化

4、裁剪不重要的层

硬件选择指南:

1、普通任务:STM32H7 系列(带 DSP/FPU)

2、高性能需求:Kendryte K210、Nordic nRF9160

3、预算充足:Raspberry Pi Zero 2W + Edge TPU

实际性能参考:

1、STM32H747:约 0.2 FPS(160×160 输入)

2、Kendryte K210:约 5 FPS(160×160 输入)

3、Raspberry Pi Zero 2W + Edge TPU:约 15 FPS(320×320 输入)

对于资源极其有限的单片机(如 Arduino Uno),建议仅处理预处理任务(如图像缩放),并将数据发送到外部设备进行推理。

相关推荐
2301_764441331 分钟前
主流手机pc品牌的端侧模型部署梳理
人工智能·windows·机器学习·智能手机·产品运营
虾壳云智能14 分钟前
阿里云百炼 API 配置 OpenClaw 2.7.9 环境搭建
人工智能·阿里云百炼·open claw安装·open claw教程
Xzh042318 分钟前
AI Agent 学习路线(Java 后端方向)
java·人工智能·学习
菜鸟的学习日记、34 分钟前
GPIO的几种模式——以STM32为例
stm32·单片机·嵌入式硬件·gpio
醒醒该学习了!35 分钟前
视觉与声音大模型(理论篇)
人工智能
Cloud_Shy61842 分钟前
解读《Effective Python 3rd Edition》:从练气到老魔(第五章 Item 33 - 35)
开发语言·人工智能·笔记·python·学习方法
救救孩子把42 分钟前
HyperFrames by HeyGen 入门教程
人工智能·视频生成·heygen
JS菌43 分钟前
AI Agent 沙箱双层防护体系:从权限过滤到内核隔离的完整实现
前端·人工智能·后端
字节跳动开源1 小时前
Viking AI 搜索 CLI—— 开发者的合法“外挂”
人工智能·agent
辰哥单片机设计1 小时前
STM32智能睡眠检测系统
stm32·单片机·嵌入式硬件