ubuntu下tensorrt加速yolov5

要在Ubuntu中使用TensorRT加速Yolov5，需要先安装Yolov5和TensorRT。

1. 安装Yolov5

可以按照官方文档的步骤进行安装：https://github.com/ultralytics/yolov5

2. 安装TensorRT

可以按照NVIDIA的官方文档进行安装：https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html

3. 导出ONNX模型

将训练好的Yolov5模型导出为ONNX格式，可以使用官方提供的导出脚本进行导出：

python export.py --weights path/to/weights.pt --img 640 --batch 1 --name yolov5s --simplify --dynamic

4. 转换为TensorRT模型

使用TensorRT提供的工具`trtexec`将ONNX模型转换为TensorRT模型：

trtexec --onnx=path/to/yolov5s.onnx --saveEngine=path/to/yolov5s.engine --explicitBatch --fp16

这里使用了`--explicitBatch`选项来指定显式批次大小，并使用了`--fp16`选项来启用FP16精度加速。

5. 加载TensorRT模型并推理

在Python代码中加载TensorRT模型，并进行推理：

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import cv2

# 加载TensorRT引擎
with open('path/to/yolov5s.engine', 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime:
    engine = runtime.deserialize_cuda_engine(f.read())

# 创建执行上下文
context = engine.create_execution_context()

# 分配输入和输出内存
input_shape = engine.get_binding_shape(0)
output_shape = engine.get_binding_shape(1)
input_host_mem = cuda.pagelocked_empty(trt.volume(input_shape), dtype=np.float32)
output_host_mem = cuda.pagelocked_empty(trt.volume(output_shape), dtype=np.float32)
input_device_mem = cuda.mem_alloc(input_host_mem.nbytes)
output_device_mem = cuda.mem_alloc(output_host_mem.nbytes)

# 加载图像并预处理
img = cv2.imread('path/to/image.jpg')
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = cv2.resize(img, (input_shape[3], input_shape[2]))
img = img.astype(np.float32) / 255.0
img = np.transpose(img, [2, 0, 1]) # HWC -> CHW
img = np.expand_dims(img, axis=0) # add batch dimension

# 将输入数据复制到设备内存
cuda.memcpy_htod(input_device_mem, input_host_mem)

# 执行推理
context.execute_v2(bindings=[int(input_device_mem), int(output_device_mem)])

# 将输出数据复制到主机内存
cuda.memcpy_dtoh(output_host_mem, output_device_mem)

# 解析输出结果
output = output_host_mem.reshape(output_shape)

这里需要注意的是，输入和输出内存的分配需要使用`cuda.pagelocked_empty`函数分配固定大小的内存，并且需要使用`cuda.mem_alloc`函数在设备上分配对应大小的内存，这样能够提高数据传输效率。

参考链接：
- https://developer.nvidia.com/blog/speeding-up-deep-learning-inference-using-tensorrt-integration-with-onnx-runtime/
- https://docs.nvidia.com/deeplearning/tensorrt/sample-support-guide/index.html#python_sample_code