yolov5 PyTorch模型转TensorRT

将YOLOv5 PyTorch模型转换为TensorRT模型需要以下步骤：

1. 安装TensorRT和PyTorch。

2. 下载并安装yolov5。

3. 使用PyTorch将yolov5模型转换为ONNX格式。

python models/export.py --weights yolov5s.pt --img 640 --batch 1 --include onnx  # yolov5s

4. 安装ONNX-TensorRT。

git clone https://github.com/onnx/onnx-tensorrt.git
cd onnx-tensorrt
git submodule update --init --recursive
mkdir build && cd build
cmake .. -DTENSORRT_ROOT=/path/to/tensorrt -DCMAKE_CXX_COMPILER=g++-7
make -j
sudo make install

5. 使用ONNX-TensorRT将ONNX模型转换为TensorRT模型。

import onnx
import onnx_tensorrt.backend as backend

model = onnx.load("yolov5s.onnx")  # Load the ONNX model
engine = backend.prepare(model, device="CUDA:0")  # Prepare the TensorRT model
with open("yolov5s.engine", "wb") as f:  # Serialize the TensorRT engine
    f.write(engine.serialize())

6. 测试TensorRT模型的性能和准确性。

import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import time

# Load the TensorRT engine
with open("yolov5s.engine", "rb") as f:
    engine = cuda.Context().deserialize_cuda_engine(f.read())

# Create the TensorRT inference context
context = engine.create_execution_context()

# Allocate the input and output buffers
input_shape = engine.get_binding_shape(0)
output_shape = engine.get_binding_shape(1)
input_buffer = cuda.mem_alloc(np.prod(input_shape) * 4)
output_buffer = cuda.mem_alloc(np.prod(output_shape) * 4)

# Prepare the input data
input_data = np.random.rand(*input_shape).astype(np.float32)

# Copy the input data to the input buffer
cuda.memcpy_htod(input_buffer, input_data)

# Run inference
start_time = time.time()
context.execute_v2(bindings=[int(input_buffer), int(output_buffer)])
end_time = time.time()

# Copy the output data to the output buffer
output_data = np.empty(output_shape, dtype=np.float32)
cuda.memcpy_dtoh(output_data, output_buffer)

# Print the inference time and output data
print("Inference time: {} ms".format((end_time - start_time) * 1000))
print("Output shape: {}".format(output_shape))
print("Output data: {}".format(output_data))