用YOLOv5和TensorRT在Jetson Nano上实现实时目标检测：一个垃圾分类项目的实战

基于YOLOv5与TensorRT的嵌入式智能垃圾分类系统实战指南当清晨第一缕阳光透过窗户洒进厨房你是否曾为这个塑料瓶该扔进哪个垃圾桶而犹豫在环保意识日益增强的今天精准垃圾分类已成为现代生活的必修课。本文将带你走进AI赋能的智能垃圾分类世界通过YOLOv5模型与TensorRT加速技术在信用卡大小的Jetson Nano开发板上构建实时识别系统——这不仅是技术实践更是对可持续生活方式的探索。1. 项目架构与技术选型1.1 为什么选择YOLOv5TensorRT组合在边缘计算场景中我们面临着三重挑战实时性要求30FPS、有限的计算资源Jetson Nano仅4核CPU128核GPU以及能耗约束典型功耗5-10W。经过对比测试YOLOv5s模型在COCO数据集上达到27.2mAP的同时仅需7.3G FLOPs其优势体现在模型效率相比YOLOv3推理速度提升2.5倍部署友好原生支持PyTorch导出ONNX格式精度平衡在自定义数据集上可达85%识别准确率TensorRT作为NVIDIA专属推理加速器通过层融合、精度校准等技术可使YOLOv5在Jetson Nano上的推理速度从12FPS提升至28FPS。典型优化策略包括优化技术效果提升实现方式FP16量化速度↑40%添加--fp16标志动态批处理吞吐量↑3x配置max_batch_size内核自动调优延迟↓15%启用tactic sources1.2 硬件配置方案针对垃圾分类场景的特殊需求我们推荐以下硬件组合# 典型硬件配置代码示例 hardware_config { 核心板: Jetson Nano 4GB B01, 摄像头: Raspberry Pi HQ Camera (1200万像素), 机械臂: SG90伺服电机x4 Arduino控制板, 扩展接口: USB3.0 x4 GPIO 40pin, 电源方案: 5V4A DC供电散热风扇 }注意避免使用Micro USB供电方式持续高负载可能导致电压不稳引发系统崩溃。建议选用官方推荐的5V4A电源适配器。2. 数据工程实战2.1 构建垃圾分类数据集不同于通用目标检测垃圾分类需要针对特定场景优化。我们在某社区收集的5类垃圾数据统计如下类别训练集验证集特点可回收物1,200300多金属反光材质厨余垃圾950250有机质变形体有害垃圾800200小目标居多其他垃圾1,500350形状不规则电子废弃物700150纹理复杂数据增强策略采用# data/augmentation.yaml hsv_h: 0.015 # 色相扰动 hsv_s: 0.7 # 饱和度增强 hsv_v: 0.4 # 明度调整 flipud: 0.5 # 垂直翻转概率 mosaic: 1.0 # 马赛克增强 mixup: 0.2 # 图像混合比例2.2 标注与格式转换使用LabelImg进行标注时需特别注意对透明材质如玻璃瓶标注外轮廓而非可见部分重叠物体采用最小外接矩形标注小目标32x32像素建议放大后标注转换YOLO格式到COCO格式的示例代码from pycocotools.coco import COCO import json def yolo_to_coco(yolo_annotations, image_size(640,640)): coco_output { images: [], annotations: [], categories: [{id: i, name: c} for i,c in enumerate(classes)] } for img_id, anns in enumerate(yolo_annotations): coco_output[images].append({ id: img_id, width: image_size[0], height: image_size[1], file_name: f{img_id}.jpg }) for ann in anns: x_center, y_center, w, h ann[bbox] coco_output[annotations].append({ id: len(coco_output[annotations]), image_id: img_id, category_id: ann[category_id], bbox: [ (x_center - w/2) * image_size[0], # x_min (y_center - h/2) * image_size[1], # y_min w * image_size[0], # width h * image_size[1] # height ], area: (w * image_size[0]) * (h * image_size[1]), iscrowd: 0 }) return coco_output3. 模型训练与优化3.1 跨设备训练方案考虑到Jetson Nano的训练性能限制建议采用云训练边缘部署模式云端训练使用AWS p3.2xlarge实例配备V100 GPUpython train.py --img 640 --batch 32 --epochs 100 --data garbage.yaml --cfg models/yolov5s.yaml --weights yolov5s.pt边缘调优在Nano上进行最后5个epoch的微调python train.py --img 320 --batch 8 --epochs 5 --device 0 --data garbage.yaml --weights yolov5s.pt关键训练参数对比参数云端值边缘值调整原因图像尺寸640320内存限制批大小328GPU显存限制学习率0.010.001防止过拟合3.2 模型剪枝与量化为提升边缘端推理效率采用以下优化组合通道剪枝from torch.nn.utils import prune model torch.load(best.pt) parameters_to_prune [(module, weight) for module in filter(lambda m: type(m) nn.Conv2d, model.modules())] prune.global_unstructured(parameters_to_prune, pruning_methodprune.L1Unstructured, amount0.3)INT8量化python export.py --weights best.pt --include onnx --int8 --img 320优化前后性能对比指标原始模型优化后提升幅度模型大小14.5MB3.7MB74%↓推理速度22FPS31FPS41%↑mAP0.50.8530.8371.6%↓4. TensorRT部署实战4.1 环境配置关键步骤Jetson Nano环境配置需要特别注意依赖版本匹配# 安装PyCUDA的正确姿势 sudo apt-get install build-essential python3-dev wget https://pypi.python.org/packages/source/p/pycuda/pycuda-2021.1.tar.gz tar xzf pycuda-2021.1.tar.gz cd pycuda-2021.1 python3 configure.py --cuda-root/usr/local/cuda-10.2 make -j4 sudo pip3 install .提示遇到nvcc not found错误时检查~/.bashrc中CUDA路径配置export PATH/usr/local/cuda-10.2/bin${PATH::${PATH}}4.2 模型转换全流程YOLOv5模型到TensorRT引擎的转换流程PT → ONNXpython export.py --weights best.pt --include onnx --img 320 --simplifyONNX → TRTimport tensorrt as trt logger trt.Logger(trt.Logger.INFO) builder trt.Builder(logger) network builder.create_network(1 int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) parser trt.OnnxParser(network, logger) with open(best.onnx, rb) as f: parser.parse(f.read()) config builder.create_builder_config() config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 30) serialized_engine builder.build_serialized_network(network, config) with open(best.engine, wb) as f: f.write(serialized_engine)验证引擎import pycuda.driver as cuda import pycuda.autoinit with open(best.engine, rb) as f: runtime trt.Runtime(logger) engine runtime.deserialize_cuda_engine(f.read()) context engine.create_execution_context() print(fEngine包含{engine.num_bindings}个绑定)4.3 实时推理优化技巧实现高效视频流处理的三个关键点流水线并行class ProcessingPipeline: def __init__(self): self.stream cuda.Stream() self.host_input cuda.pagelocked_empty(INPUT_SIZE, dtypenp.float32) self.device_input cuda.mem_alloc(self.host_input.nbytes) def async_inference(self, frame): preprocessed self.preprocess(frame) np.copyto(self.host_input, preprocessed.ravel()) cuda.memcpy_htod_async(self.device_input, self.host_input, self.stream) context.execute_async_v2(bindings[int(self.device_input), int(self.device_output)], stream_handleself.stream.handle) cuda.memcpy_dtoh_async(self.host_output, self.device_output, self.stream) self.stream.synchronize() return self.postprocess(self.host_output)内存复用通过cuda.MemoryPool创建内存池避免频繁分配释放动态分辨率处理def set_dynamic_shape(context): profile builder.create_optimization_profile() profile.set_shape(input, (1,3,320,320), (1,3,640,640), (1,3,640,640)) config.add_optimization_profile(profile)5. 系统集成与性能调优5.1 机械臂控制联动通过GPIO控制垃圾分类机械臂的示例代码import Jetson.GPIO as GPIO import time GPIO.setmode(GPIO.BOARD) GPIO.setup(11, GPIO.OUT) # 可回收物 GPIO.setup(13, GPIO.OUT) # 厨余垃圾 GPIO.setup(15, GPIO.OUT) # 有害垃圾 def control_arm(class_id): pins {0:11, 1:13, 2:15, 3:16} GPIO.output(pins[class_id], GPIO.HIGH) time.sleep(1) # 保持动作1秒 GPIO.output(pins[class_id], GPIO.LOW)5.2 端到端延迟分析使用NVIDIA Nsight Systems进行性能剖析nsys profile -t cuda,nvtx --statstrue python detect.py --source 0典型性能瓶颈及解决方案摄像头采集延迟~50ms改用V4L2直接读取接口降低分辨率到720p后处理耗时~20ms使用CUDA实现NMS算法将分类结果绘制移到单独线程显示输出延迟~30ms禁用GUI直接输出到机械臂使用硬件加速的OpenGL渲染5.3 实际部署注意事项在社区垃圾站连续运行30天的经验总结散热处理加装散热片风扇使SoC温度稳定在65℃以下电源管理配置看门狗定时器防止系统卡死模型更新每周通过OTA更新模型权重异常处理对连续10帧检测失败自动重启服务# 看门狗服务配置示例 */5 * * * * /usr/bin/python3 /home/nano/watchdog.py在项目落地过程中最令人意外的是光照条件对识别率的影响——黄昏时分由于色温变化塑料瓶的误识别率会上升15%。我们最终通过增加HSV色彩扰动训练数据解决了这个问题。这也印证了边缘AI项目的黄金准则没有放之四海皆准的完美模型只有不断迭代优化的场景适配。