三亩地 三亩地SAN MU DI · CODE DIARY
ARTICLE DETAIL

日记详情

真实记录编程学习的某一天,欢迎挑你感兴趣的翻一翻。

从YOLOv5的.pt到.bin:模型轻量化与端侧部署实战指南

从YOLOv5的.pt到.bin:模型轻量化与端侧部署实战指南

1. 为什么需要从.pt转换到.bin?

当你用YOLOv5训练出一个效果不错的检测模型,拿到那个best.pt文件时,可能已经迫不及待想把它部署到手机或嵌入式设备上了。但这时候你会发现,这些设备根本吃不消原始的PyTorch模型。我去年给工厂做缺陷检测项目时就踩过这个坑——在服务器上跑得好好的模型,放到产线工控机上直接卡成PPT。

.bin格式其实就是经过优化的二进制模型文件,它比.pt小得多(通常能压缩到1/3大小),而且能被大多数端侧推理框架直接加载。比如在安卓上用NCNN、在树莓派上用Tengine,都需要这种轻量化格式。不过转换过程就像把一道中式炒菜改造成速冻食品,既要保持风味又要适应新场景,中间有几个关键步骤必须处理好。

2. 从.pt到ONNX:跨框架的第一步

2.1 准备工作就像备菜

首先确保你的YOLOv5环境完整,我推荐用官方最新的v6.1版本。遇到过有人用老版本转换后输出维度对不齐的情况,就像炒菜时发现酱油过期了。关键依赖要装全:

pip install onnx coremltools opencv-python

你的.pt文件最好是直接训练输出的(比如best.pt),不要自己手动改过结构。我有次手贱在模型里加了自定义层,转换时直接报错,就像非要把土豆塞进豆浆机。

2.2 执行转换命令的玄机

在yolov5目录下运行这个命令:

python export.py --weights best.pt --img 640 --batch 1 --include onnx

注意几个参数陷阱:

  • --img 640必须和训练时一致,我试过改成320,结果检测框全飘了
  • --batch 1是固定写法,端侧推理都是单张处理
  • 新版已经不需要手动指定--train,系统会自动处理

转换完成后用Netron打开onnx文件,应该能看到三个输出节点(如下图)。如果发现输出维度是"?"而不是具体数字,说明转换时漏了动态维度设置,后面部署会出大问题。

3. ONNX瘦身:给模型做抽脂手术

3.1 为什么简化是必须的?

原始ONNX模型包含大量训练用的冗余节点,就像带着全套露营装备去上班。直接转换会:

  1. 增加20%-30%的无用体积
  2. 可能包含不支持的算子
  3. 影响后续优化器效果

3.2 实操中的隐藏技巧

安装简化工具:

pip install onnx-simplifier

运行这个命令时要特别注意:

python -m onnxsim best.onnx best-sim.onnx --input-shape 1,3,640,640

那个--input-shape参数就像给模型定制的紧身衣,能帮优化器更好地裁剪。有次我忘了加这个,转换后的模型在安卓上输入尺寸死活对不上。

4. 终极转换:生成.bin部署文件

4.1 在线转换工具的妙用

推荐用这个神器网站:convertmodel.com(注意:国内访问可能需要特殊网络配置)。操作步骤:

  1. 上传刚才的best-sim.onnx
  2. 选择输出格式为NCNN
  3. 点击Convert后下载.zip包

解压后会得到两个文件:

  • .param:模型结构描述
  • .bin:模型权重数据

4.2 手动调参的生存指南

用文本编辑器打开.param文件,找到最后的输出层(通常是三个DetectionOutput节点),把里面的数字全改成-1。这个操作就像给模型装上安全阀,防止安卓端出现检测框爆炸的情况。

我整理了个对比表格方便理解:

参数位置修改前修改后作用
output125200-1避免密集检测
output225200-1同上
output325200-1同上

5. 安卓端部署实战心得

5.1 环境配置的坑

在Android Studio里引入NCNN时,要注意这些配置:

android { defaultConfig { ndk { abiFilters 'armeabi-v7a', 'arm64-v8a' } } }

有次我漏了abiFilters,在小米手机上跑得好好的,到华为就直接闪退,查了三天才发现是架构兼容问题。

5.2 推理代码的优化技巧

加载模型时要特别注意路径问题:

ncnn.Net net = new ncnn.Net(); net.loadParam("models/yolo-opt.param"); net.loadModel("models/yolo-opt.bin");

实测发现,把模型文件放在assets目录会比直接放res快30%左右。更骚的操作是预加载到内存,启动速度能提升近一倍。

6. 性能调优的那些事儿

模型转换只是开始,真正的挑战在于让它在设备上飞起来。分享几个压箱底的技巧:

  1. 量化到8位:用NCNN的quantize工具,速度能提升3倍
  2. 多线程推理:设置net.set_num_threads(4)充分利用多核
  3. 输入尺寸优化:不一定非要用640x640,试试480x480可能精度只降1%但速度快50%

最近给某家电厂商做的电饭煲缺陷检测项目,经过这些优化后,在联发科低端芯片上都能跑到35FPS,比他们原来用的云端方案快得多。

← 返回列表