
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 = µ_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 转换为适合单片机的格式
资源限制与优化建议
- 模型选择:优先选用 Nano 或 Tiny 版本的模型
- 输入尺寸:使用较小的输入尺寸(如 160×160 或 320×320)
- 量化:使用 8 位或 16 位量化,甚至二值化
- 层融合:利用框架的层融合功能减少计算量
- 算法简化:只检测关键类别,降低模型复杂度
对于资源极其有限的单片机(如 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),建议仅处理预处理任务(如图像缩放),并将数据发送到外部设备进行推理。