官方yolo v5运行.onnx脚本

### 将YOLOv5模型转换为ONNX格式 为了将YOLOv5模型转换成ONNX格式，可以利用PyTorch中的`torch.onnx.export()`函数来实现这一目标[^1]。具体来说，在Python环境中执行此操作时，需要先加载训练好的YOLOv5模型实例，并指定输入张量的形状以及导出文件的目标路径。 以下是具体的代码示例： ```python import torch from models.experimental import attempt_load device = 'cuda' if torch.cuda.is_available() else 'cpu' model = attempt_load('yolov5s.pt', map_location=device) # 加载预训练权重 model.eval() dummy_input = torch.randn(1, 3, 640, 640).to(device) output_onnx = 'yolov5s.onnx' input_names = ["images"] output_names = ['output'] dynamic_axes = {'images': {0: 'batch_size'}, 'output': {0: 'batch_size'}} torch.onnx.export(model, dummy_input, output_onnx, input_names=input_names, dynamic_axes=dynamic_axes, opset_version=12) ``` 这段脚本会创建一个名为'yolov5s.onnx'的新文件，该文件即为转换后的ONNX模型表示形式。通过设置动态轴参数，使得批处理大小能够灵活调整；而opset版本号则决定了所使用的ONNX运算符集版本。 一旦完成了上述过程之后，就可以借助像ONNX Runtime这样的工具来进行推理工作了。对于想要进一步优化性能的情况，则可能涉及到量化等高级特性[^2]。

### YOLO v5 教程与使用 YOLO (You Only Look Once) 是一种实时目标检测算法，而 YOLO v5 则是该系列中的一个重要版本。它提供了友好的架构设计和高效的性能表现[^2]。 #### 安装依赖 要开始使用 YOLO v5，首先需要安装必要的 Python 库和其他环境配置。以下是基本的安装步骤： 1. **克隆仓库** 使用 Git 将官方 YOLO v5 仓库克隆到本地： ```bash git clone https://github.com/ultralytics/yolov5.git cd yolov5 ``` 2. **创建虚拟环境并安装依赖项** 创建一个新的虚拟环境，并安装 `requirements.txt` 中列出的所有依赖库： ```bash python3 -m venv env source env/bin/activate pip install -r requirements.txt ``` 3. **验证安装成功** 运行以下命令以测试安装是否正常工作： ```bash python detect.py --weights yolov5s.pt --img 640 --conf 0.25 --source data/images/ ``` 如果一切顺利，则会在终端中显示推理结果，并保存预测图像至指定目录。 --- #### TensorRT 部署流程 对于高性能需求场景下部署 YOLO v5 模型时可以考虑 NVIDIA 的 TensorRT 工具包。具体实现如下所示： 1. **导出 ONNX 文件** 转换 PyTorch 权重文件为通用中间表示形式 `.onnx`。 ```python import torch model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5s.pt') dummy_input = torch.randn(1, 3, 640, 640) torch.onnx.export( model, dummy_input, "yolov5s.onnx", opset_version=12, do_constant_folding=True, input_names=['input'], output_names=['output'] ) ``` 2. **加载 ONNX 并构建 TensorRT 引擎** 下面展示如何利用 C++ API 加载已有的 ONNX 文件并生成优化后的引擎文件。 ```cpp #include <NvInfer.h> #include <fstream> using namespace nvinfer1; IBuilder* builder = createInferBuilder(gLogger); INetworkDefinition* network = builder->createNetworkV2(0U); // 解析器初始化 IParse* parser = nvonnxparser::createParser(*network, gLogger); std::ifstream file("yolov5s.onnx", std::ios::binary | std::ios::ate); size_t fileSize = static_cast<size_t>(file.tellg()); char *buffer = new char[fileSize]; file.seekg(0, std::ios::beg).read(buffer, fileSize); parser->parseFromBuffer(buffer, fileSize); delete[] buffer; parser->destroy(); IBuilderConfig* config = builder->createBuilderConfig(); ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); assert(engine != nullptr && "Failed to build the engine"); serializeAndSaveEngine(engine, "yolov5s.trt"); destroyResources(network, config, builder); ``` 以上代码片段展示了从 ONNX 导入模型直至最终序列化存储整个过程[^1]。 --- #### 训练自定义数据集 如果想基于自己的标注图片微调预训练权重或者重新训练新类别识别能力的话，可以通过修改配置文件完成相应调整操作。 1. 准备好 COCO 或者 VOC 类型的数据结构； 2. 编辑 YAML 格式的参数说明文档（如路径、数量统计等信息）； 3. 执行脚本启动学习进程： ```bash python train.py --img 640 --batch 16 --epochs 50 \ --data custom_dataset.yaml --cfg models/yolov5s.yaml \ --weights '' --name yolo_custom_run_1 ``` 更多细节可查阅源码解读资料[^3]。 ---