YOLOv8极速CPU优化:物联网设备毫秒级推理的代码实现与性能调优
做边缘物联网设备落地的同学肯定都遇到过这个痛点:很多低功耗设备(门禁摄像头、智能家居网关、老旧工业设备)只有低端CPU,没有GPU甚至NPU,原生YOLOv8n跑640×640输入要300ms以上,根本达不到实时要求,换硬件成本又太高,项目根本没法落地。
去年我做老旧小区门禁改造项目,现场的门禁设备只有ARM Cortex-A53 4核CPU,主频1.5GHz,没有任何加速单元,要求人脸和人员检测达到30FPS,也就是单帧推理≤33ms。我用了5层优化手段,把原来320ms的推理速度优化到了28ms,精度只掉了0.6%,完全满足项目要求,不用换硬件,软件升级就搞定了,甲方特别满意。
今天把完整的CPU优化方案分享给大家,不管是ARM还是X86的CPU,照着做至少能把YOLO的推理速度提升5倍以上,纯CPU也能跑到实时。
一、优化前后效果对比
先放一下我在ARM Cortex-A53 CPU上的优化效果,输入分辨率640×640,模型是YOLOv8n:
| 优化阶段 | 推理耗时 | FPS | mAP@0.5 | 优化幅度 |
|---|---|---|---|---|
| 原生PyTorch推理 | 320ms | 3.1 | 92.3% | - |
| 导出ONNX+ONNX Runtime CPU推理 | 185ms | 5.4 | 92.3% | 提升73% |
| 结构化剪枝+知识蒸馏 | 126ms | 7.9 | 91.9% | 提升47% |
| 算子融合+图优化 | 78ms | 12.8 | 91.9% | 提升62% |
| INT8量化 | 38ms | 26.3 | 91.7% | 提升105% |
| 预处理+后处理优化 | 28ms | 35.7 | 91.7% | 提升36% |
最终耗时28ms,比原生快了11倍,精度只掉了0.6%,完全满足实时要求。
二、逐层优化代码实现
2.1 第一层:ONNX导出+运行时优化
这是最基础的优化,提升效果最明显,而且完全不损失精度,5分钟就能搞定。
首先导出优化后的ONNX模型:
# 导出ONNX,带上算子融合,opset选17,对CPU优化更好yoloexportmodel=yolov8n.ptformat=onnxsimplify=Trueopset=17dynamic=Falsesimplify=True会自动做第一轮的算子融合,删除无用的算子,模型大小会减少20%左右。
然后用ONNX Runtime的CPU优化版本推理,不要用PyTorch推理,速度差好几倍:
importonnxruntimeasrtimportnumpyasnpimportcv2# 配置ONNX Runtime CPU优化选项session_options=rt.SessionOptions()# 开启图优化session_options.graph_optimization_level=rt.GraphOptimizationLevel.ORT_ENABLE_ALL# 开启线程优化,设置CPU核心数,ARM A53是4核,设为4session_options.intra_op_num_threads=4session_options.inter_op_num_threads=1# 开启内存优化session_options.enable_mem_pattern=Truesession_options.enable_cpu_mem_arena=True# 加载模型,用CPU执行提供者session=rt.InferenceSession("yolov8n.onnx",sess_options=session_options,providers=["CPUExecutionProvider"])这一步做完,速度直接提升70%以上,而且精度完全没有损失。如果是X86的CPU,推荐用OpenVINO执行提供者,速度比ONNX Runtime还要快20%左右:
providers=["OpenVINOExecutionProvider"]2.2 第二层:结构化剪枝+知识蒸馏
剪枝可以把模型里没用的通道剪掉,减少参数量和计算量,速度提升很多。我用的是TorchPrune的结构化剪枝,剪枝比例40%,然后用知识蒸馏恢复精度:
importtorchimporttorch_pruningastpfromultralyticsimportYOLO# 加载教师模型和学生模型teacher_model=YOLO("yolov8n.pt").model student_model=YOLO("yolov8n.pt").model# 剪枝配置example_inputs=torch.randn(1,3,640,640)ignored_layers=[student_model.model[-1]]# 不要剪检测头pruner=tp.pruner.MagnitudePruner(student_model,example_inputs=example_inputs,importance=tp.importance.GroupNormImportance(p=2),global_pruning=False,pruning_ratio=0.4,# 剪枝40%的通道ignored_layers=ignored_layers)pruner.step()# 剪枝后用知识蒸馏训练15个epoch恢复精度# 蒸馏的loss是硬标签loss+软标签loss的加权和defdistillation_loss(student_outputs,teacher_outputs,targets,alpha=0.3,temperature=4):hard_loss=student_model.loss(student_outputs,targets)# 软标签蒸馏losssoft_loss=nn.KLDivLoss(reduction="batchmean")(F.log_softmax(student_outputs[0]/temperature,dim=1),F.softmax(teacher_outputs[0]/temperature,dim=1))*(temperature**2)returnalpha*hard_loss+(1-alpha)*soft_loss# 训练过程省略,只需要训练15个epoch,学习率用原来的1/10剪枝之后参数量减少40%,速度提升47%,经过知识蒸馏之后精度只掉了0.4%,几乎可以忽略。
2.3 第三层:算子融合+图优化
这一步用ONNX Runtime的工具做离线的图优化,把多个小算子融合成大算子,减少内存拷贝和kernel调用的开销:
# 用onnxruntime-tools优化ONNX模型python-monnxruntime.quantization.preprocess--inputyolov8n-pruned.onnx--outputyolov8n-pruned-opt.onnx这个工具会自动做:
- 卷积+BN+激活融合
- 矩阵乘法+偏置融合
- 冗余算子删除
- 常量折叠
- 内存布局优化
优化之后模型的计算量减少20%左右,速度再提升60%,精度完全没有损失。
2.4 第四层:INT8量化
INT8量化是CPU优化的大杀器,把32位浮点运算变成8位整数运算,速度可以提升一倍,精度损失控制得好的话可以做到小于1%。
我用的是ONNX Runtime的静态量化,需要用校准集校准:
fromonnxruntime.quantizationimportquantize_static,CalibrationDataReader,QuantType# 校准数据读取器,用100-200张和实际场景一致的图片做校准classMyCalibrationDataReader(CalibrationDataReader):def__init__(self,image_dir):self.image_list=[os.path.join(image_dir,f)forfinos.listdir(image_dir)iff.endswith(".jpg")]self.index=0self.input_name="images"defget_next(self):ifself.index>=len(self.image_list):returnNoneimg=cv2.imread(self.image_list[self.index])img=cv2.resize(img,(640,640))img=img.transpose(2,0,1)[np.newaxis,:,:,:].astype(np.float32)/255.0self.index+=1return{self.input_name:img}# 静态量化quantize_static(model_input="yolov8n-pruned-opt.onnx",model_output="yolov8n-int8.onnx",calibration_data_reader=MyCalibrationDataReader("datasets/calib/"),quant_format=QuantType.QInt8,op_types_to_quantize=["Conv","MatMul"],# 只量化卷积和矩阵乘法,其他算子保持FP32,减少精度损失per_channel=True,reduce_range=True# ARM CPU推荐打开reduce_range,精度更好)量化之后速度直接提升一倍,从78ms降到38ms,精度只掉了0.2%,几乎感知不到。注意校准集一定要选和实际部署场景分布一致的图片,不然精度会掉很多,我用了200张实际门禁场景的图片做校准,精度只掉了0.2%。
2.5 第五层:预处理+后处理优化
很多人忽略了预处理和后处理的耗时,其实这部分经常能占到总耗时的30%以上,优化空间很大。
预处理优化:
原来的预处理是用Python写的,循环+numpy操作,很慢,优化成OpenCV的向量化操作,或者用C++写预处理,速度提升特别明显:
# 优化前:numpy操作,耗时约8msdefpreprocess_slow(img):img=cv2.resize(img,(640,640))img=img[:,:,::-1].transpose(2,0,1)# BGR转RGB,HWC转CHWimg=np.ascontiguousarray(img,dtype=np.float32)/255.0img=np.expand_dims(img,axis=0)returnimg# 优化后:OpenCV向量化操作,耗时约2msdefpreprocess_fast(img):h,w=img.shape[:2]scale=min(640/h,640/w)new_h,new_w=int(h*scale),int(w*scale)# 用cv2.INTER_LINEAR_EXACT更快,精度差不多img_resized=cv2.resize(img,(new_w,new_h),interpolation=cv2.INTER_LINEAR_EXACT)# 直接用cv2.cvtColor转RGB,比numpy索引快img_rgb=cv2.cvtColor(img_resized,cv2.COLOR_BGR2RGB)# 用cv2.copyMakeBorder填充,比numpy快pad_h,pad_w=(640-new_h)//2,(640-new_w)//2img_padded=cv2.copyMakeBorder(img_rgb,pad_h,640-new_h-pad_h,pad_w,640-new_w-pad_w,cv2.BORDER_CONSTANT,value=(114,114,114))# 归一化用cv2.convertScaleAbs更快,或者直接在推理的时候做融合img_input=img_padded.transpose(2,0,1)[np.newaxis,:,:,:].astype(np.float32)/255.0returnimg_input预处理耗时从8ms降到2ms。
后处理优化:
把后处理里的循环操作改成numpy向量化操作,或者用Cython编译,耗时从10ms降到3ms。
这样预处理+后处理总共耗时从18ms降到5ms,总耗时从38ms降到28ms,达到35FPS。
三、不同CPU的优化效果
我在几种常见的CPU上都做了测试,输入都是640×640,优化后的速度如下:
| CPU型号 | 核心 | 优化前耗时 | 优化后耗时 | FPS |
|---|---|---|---|---|
| ARM Cortex-A53 1.5GHz | 4核 | 320ms | 28ms | 35.7 |
| ARM Cortex-A76 2.0GHz | 8核 | 120ms | 11ms | 90.9 |
| Intel i5-10400 | 6核 | 80ms | 6ms | 166.7 |
| Intel Celeron J4125 | 4核 | 180ms | 16ms | 62.5 |
不管是低端ARM还是X86 CPU,优化之后都能达到实时的速度,完全满足大多数物联网场景的需求。
四、落地避坑指南
- 校准集不要用公开数据集:一定要用实际部署场景的图片做校准,我最开始用COCO数据集做校准,量化之后精度掉了5%,换成实际场景的200张图片之后,精度只掉了0.2%
- ARM CPU选对量化参数:ARM CPU的INT8运算和X86不一样,一定要打开
reduce_range=True,不然精度会掉很多 - 线程数不要设太大:对于4核CPU,intra_op_num_threads设为4就够了,设太大反而会因为线程调度开销导致速度变慢
- 不要盲目追求高剪枝比例:剪枝比例太高的话,就算蒸馏也恢复不了精度,40%-50%是比较合适的比例,最多不要超过60%
- 输入分辨率不要盲目降:很多人为了速度把输入分辨率降到320×320,精度掉的特别多,通过上面的优化手段,640×640也能跑到实时,精度高很多
这套CPU优化方案我已经在十几个物联网项目里落地了,包括门禁、智能家居、工业检测等场景,不需要换硬件,纯软件优化就能达到实时要求,大大降低了项目落地的成本,大家有物联网设备部署需求的一定要试试。
