ubuntu yolov5 tensorRT

### 将 YOLOv5 转换为 TensorRT 格式的流程 为了在 Ubuntu 上将 YOLOv5 模型转换为 TensorRT 格式以加速推理性能，可以按照以下方法操作： #### 1. 安装必要的依赖项 确保安装了 OpenCV 和其他必要工具[^1]。此外，还需要安装 NVIDIA 的 TensorRT 库以及 PyTorch 版本兼容的支持包。 ```bash sudo apt update && sudo apt install libopencv-dev python3-opencv pip install torch torchvision numpy onnx ``` 对于 TensorRT 的安装，可以通过官方文档获取支持版本的安装命令。例如： ```bash # 替换 <version> 为你需要的具体版本号 pip install nvidia-pyindex pip install nvidia-tensorrt==<version> ``` #### 2. 导出 ONNX 文件 YOLOv5 支持导出为多种格式，其中包括 ONNX (Open Neural Network Exchange)，这是 TensorRT 所需的第一步中间文件。 运行以下脚本来导出模型为 ONNX 格式： ```python import torch model = torch.hub.load('ultralytics/yolov5', 'custom', path='path/to/best.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'] ) print("ONNX export complete.") ``` 此过程会生成 `yolov5s.onnx` 文件作为下一步输入。 #### 3. 使用 TensorRT 进行优化 通过 Python 或 C++ 接口加载 ONNX 并将其转化为 TensorRT 引擎文件。以下是基于 Python API 的实现方式： ```python import tensorrt as trt TRT_LOGGER = trt.Logger(trt.Logger.WARNING) def build_engine(onnx_file_path): with trt.Builder(TRT_LOGGER) as builder, \ builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) as network, \ trt.OnnxParser(network, TRT_LOGGER) as parser: builder.max_workspace_size = 1 << 30 # 设置最大工作空间大小为 1GB builder.max_batch_size = 1 # 设定批量大小 with open(onnx_file_path, 'rb') as model: if not parser.parse(model.read()): print('Failed to parse the ONNX file.') for error in range(parser.num_errors): print(parser.get_error(error)) return None engine = builder.build_cuda_engine(network) return engine engine = build_engine("yolov5s.onnx") with open("yolov5s.trt", "wb") as f: f.write(engine.serialize()) print("TensorRT Engine saved successfully.") ``` 上述代码完成从 `.onnx` 到 `.trt` 工程化的过程[^2]。 #### 4. 验证和测试 最后一步是对生成的 TensorRT 引擎进行验证，确认其推理速度提升效果显著的同时保持精度不变。 ---