yolo onnx 树莓派

### 将YOLO模型转换为ONNX格式并在树莓派上进行推理 #### 准备工作 为了在树莓派上成功将YOLO模型转换成ONNX格式并执行推理操作，需先安装必要的依赖库和环境配置。 对于Python环境而言，建议创建虚拟环境来管理项目所需的包。可以利用`venv`模块建立新的虚拟环境： ```bash python3 -m venv yolov5-env source yolov5-env/bin/activate ``` 接着，在激活后的环境中通过pip命令安装PyTorch和其他必需的软件包[^1]。 #### 安装依赖项 确保已经安装了最新版的pytorch, torchvision 和 onnx: ```bash pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu pip install onnx==1.12.0 onnxruntime-gpu opencv-python matplotlib scikit-image ``` #### 转换YOLO模型至ONNX格式 使用官方提供的脚本或自定义代码片段可轻松地把`.pt`(PyTorch)文件转化为`.onnx`文件形式。这里给出一段简单的Python代码用于展示这一过程： ```python import torch.onnx as onnx from models.experimental import attempt_load import torch weights_path = "path/to/best.pt" model = attempt_load(weights=weights_path, map_location=torch.device('cpu')) dummy_input = torch.randn(1, 3, 640, 640) output_onnx = 'best_model.onnx' input_names = ["images"] output_names = ['num_detections', 'boxes', 'scores', 'classes'] dynamic_axes = { input_names[0]: {0: 'batch_size'}, } torch.onnx.export( model, dummy_input, output_onnx, verbose=False, export_params=True, do_constant_folding=True, input_names=input_names, output_names=output_names, dynamic_axes=dynamic_axes, opset_version=11 ) print(f"Model has been converted to ONNX and saved at {output_onnx}") ``` 这段代码会读取预训练好的YOLO权重文件(`best.pt`)，并通过给定的一个随机张量作为输入样例来进行导出操作，最终保存为名为`best_model.onnx`的新文件。 #### 配置OnnxRuntime以支持硬件加速 为了让ONNX Runtime能够在ARM架构下的设备比如树莓派上有更好的性能表现，应该下载适用于Linux ARM64平台上的特定版本，并按照文档说明完成设置。这一步骤能够显著提升后续推断的速度与效率[^2]。 #### 运行推理测试 一旦完成了上述准备工作之后就可以编写一小段程序用来加载刚刚生成出来的ONNX模型并对图片实施检测任务了。下面是一个基本的例子： ```python import cv2 import numpy as np import onnxruntime as ort def preprocess(image): img = cv2.resize(image, (640, 640)) img = img.transpose((2, 0, 1))[::-1] img = np.ascontiguousarray(img).astype(np.float32)/255. return img[np.newaxis,...] ort_session = ort.InferenceSession("best_model.onnx") image = cv2.imread("test.jpg") # 替换成自己的图像路径 inputs = {'images': preprocess(image)} outputs = ort_session.run(None, inputs)[0] for det in outputs: box = det[:4].round().astype(int) conf = det[4] cls = int(det[5]) label = f'{conf:.2f}' plot_one_box(box, image, label=label, color=(0,255,0), line_thickness=3) cv2.imshow("Result", image) cv2.waitKey(0) ``` 此部分展示了如何调用ONNX Runtime API去处理实际的数据样本，同时包含了对预测结果可视化的一些简单逻辑。