<think>我们正在使用CUDA 12.6和TensorRT 10.9，这是一个较新的组合。由于YOLOv5官方可能尚未正式支持TensorRT 10.9，因此我们需要确保使用兼容的YOLOv5版本和正确的导出方式。
根据之前的引用[2]，我们知道在TensorRT 8.6.1上导出YOLOv5模型为engine的步骤。现在我们需要适配到TensorRT 10.9。
步骤：
1. 导出模型：将YOLOv5的PyTorch模型（.pt）导出为ONNX格式，然后再转换为TensorRT引擎（.engine）。但是，TensorRT 10.9可能已经直接支持从PyTorch模型转换，但为了通用性，我们还是采用ONNX作为中间格式。
2. 使用TensorRT Python API加载引擎并进行推理。
注意：由于TensorRT 10.9较新，我们需要确保使用的PyTorch版本和ONNX版本与之兼容。
参考步骤：
- 首先，确保环境：CUDA 12.6, cuDNN（建议与CUDA 12.6匹配的版本），TensorRT 10.9，以及对应的PyTorch（支持CUDA 12.6的版本，如PyTorch 2.3+）。
- 导出ONNX：使用YOLOv5的export.py脚本，指定opset（建议使用13或更高，因为TensorRT 10.9支持较高的ONNX opset）。
- 使用TensorRT的trtexec工具或者TensorRT Python API将ONNX转换为engine。这里我们使用Python API，以便在代码中完成转换和推理。
但是，用户要求直接编写测试YOLOv5模型推理的代码，所以我们假设已经有一个转换好的TensorRT引擎文件（.engine）。
因此，代码步骤：
1. 导入必要的库：tensorrt, pycuda, numpy, opencv等。
2. 加载TensorRT引擎文件。
3. 创建执行上下文。
4. 准备输入数据（图像预处理）。
5. 执行推理。
6. 后处理（解析检测结果，应用NMS等）。
7. 绘制检测框并显示。
注意：由于TensorRT 10.9的API可能有变化，我们参考TensorRT官方Python示例。
由于代码较长，我们将分步骤编写。这里假设已经有一个转换好的yolov5s.engine文件（针对TensorRT 10.9和CUDA 12.6转换的）。
但是，请注意：YOLOv5的TensorRT引擎在导出时可能已经固定了输入尺寸，所以预处理需要按照导出时的尺寸进行。
我们将编写一个完整的测试脚本，包括以下部分：
- 图像预处理：调整大小、归一化、转换为CHW格式等。
- 加载引擎并分配内存（输入和输出）。
- 执行推理。
- 后处理：将输出数据转换为检测框，并应用非极大值抑制（NMS）。
- 绘制结果。
由于TensorRT 10.9的API与之前版本可能略有不同，我们将按照TensorRT 10.9的文档来写。
注意：由于TensorRT 10.9是较新的版本，我们可能需要使用最新的API。以下代码基于TensorRT的Python示例编写。
步骤：
1. 导入库
2. 加载引擎
3. 创建上下文
4. 准备输入和输出的内存
5. 图像预处理
6. 执行推理
7. 后处理
8. 可视化
由于YOLOv5的输出层结构可能因版本而异，我们假设使用YOLOv5 v6.0或更高版本，其输出为单个张量（形状为[1, num_boxes, 85]），其中85表示每个检测框的坐标（4个值）、目标得分（1个值）和80个类别的概率（80个值）。
但是，在TensorRT中，当我们导出为engine时，输出层的形状和名称可能会变化。因此，我们需要知道引擎的输入和输出名称以及形状。我们可以通过引擎的信息来获取。
由于代码较长，我们将分步骤实现。
注意：由于我们使用的是TensorRT 10.9，所以需要安装对应的Python包（tensorrt）和pycuda（用于GPU内存管理和数据传输）。
如果还没有转换引擎，我们可以先写一个函数来从ONNX转换，但用户要求的是测试推理，所以我们假设引擎已经存在。
开始编写代码：</think>以下是为 CUDA 12.6 和 TensorRT 10.9 编写的 YOLOv5 模型推理测试代码，包含完整的引擎加载、预处理、推理和后处理流程：
```python
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import cv2
import time
import os
class YOLOv5TRTInference:
def __init__(self, engine_path, input_shape=(640, 640)):
"""
初始化TensorRT推理引擎
:param engine_path: TensorRT引擎文件路径
:param input_shape: 模型输入尺寸 (高度, 宽度)
"""
self.input_shape = input_shape
self.logger = trt.Logger(trt.Logger.INFO)
# 加载TensorRT引擎
with open(engine_path, "rb") as f, trt.Runtime(self.logger) as runtime:
self.engine = runtime.deserialize_cuda_engine(f.read())
# 创建执行上下文
self.context = self.engine.create_execution_context()
# 分配输入输出内存
self.inputs, self.outputs, self.bindings, self.stream = self.allocate_buffers()
def allocate_buffers(self):
"""分配GPU内存并创建流"""
inputs, outputs, bindings = [], [], []
stream = cuda.Stream()
for binding in self.engine:
size = trt.volume(self.engine.get_binding_shape(binding))
dtype = trt.nptype(self.engine.get_binding_dtype(binding))
# 分配主机和设备内存
host_mem = cuda.pagelocked_empty(size, dtype)
device_mem = cuda.mem_alloc(host_mem.nbytes)
bindings.append(int(device_mem))
if self.engine.binding_is_input(binding):
inputs.append({'host': host_mem, 'device': device_mem})
else:
outputs.append({'host': host_mem, 'device': device_mem})
return inputs, outputs, bindings, stream
def preprocess_image(self, image_path):
"""
图像预处理
:param image_path: 输入图像路径
:return: 预处理后的图像数组
"""
# 读取并调整图像尺寸
img = cv2.imread(image_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img_resized = cv2.resize(img, (self.input_shape[1], self.input_shape[0]))
# 归一化并转换为CHW格式
img_norm = img_resized.astype(np.float32) / 255.0
img_chw = np.transpose(img_norm, (2, 0, 1))
# 添加批次维度
return np.expand_dims(img_chw, axis=0)
def inference(self, input_data):
"""执行推理"""
# 复制数据到GPU
np.copyto(self.inputs[0]['host'], input_data.ravel())
cuda.memcpy_htod_async(self.inputs[0]['device'], self.inputs[0]['host'], self.stream)
# 执行推理
self.context.execute_async_v2(
bindings=self.bindings,
stream_handle=self.stream.handle
)
# 将结果复制回主机
for out in self.outputs:
cuda.memcpy_dtoh_async(out['host'], out['device'], self.stream)
self.stream.synchronize()
# 返回所有输出
return [out['host'] for out in self.outputs]
def postprocess(self, outputs, conf_thres=0.5, iou_thres=0.45):
"""
后处理：非极大值抑制(NMS)
:param outputs: 模型输出
:param conf_thres: 置信度阈值
:param iou_thres: IoU阈值
:return: 检测结果列表 [x1, y1, x2, y2, confidence, class_id]
"""
# 假设输出为[1, num_boxes, 85]格式
predictions = np.squeeze(outputs[0])
# 过滤低置信度检测
conf_mask = predictions[:, 4] > conf_thres
predictions = predictions[conf_mask]
if len(predictions) == 0:
return []
# 获取类别分数
scores = predictions[:, 5:] * predictions[:, 4:5]
class_ids = np.argmax(scores, axis=1)
class_scores = scores[np.arange(len(scores)), class_ids]
# 获取边界框坐标
boxes = predictions[:, :4]
boxes[:, 0] -= boxes[:, 2] / 2 # x_center -> x1
boxes[:, 1] -= boxes[:, 3] / 2 # y_center -> y1
# 执行NMS
indices = cv2.dnn.NMSBoxes(
boxes.tolist(),
class_scores.tolist(),
conf_thres,
iou_thres
)
if len(indices) > 0:
indices = indices.flatten()
return np.hstack((
boxes[indices],
class_scores[indices].reshape(-1, 1),
class_ids[indices].reshape(-1, 1)
))
return []
def visualize(self, image_path, detections):
"""可视化检测结果"""
img = cv2.imread(image_path)
h, w = img.shape[:2]
# 缩放因子
scale_x = w / self.input_shape[1]
scale_y = h / self.input_shape[0]
for det in detections:
x1, y1, x2, y2, conf, cls_id = det
# 缩放回原始尺寸
x1 = int(x1 * scale_x)
y1 = int(y1 * scale_y)
x2 = int((x1 + x2) * scale_x)
y2 = int((y1 + y2) * scale_y)
# 绘制边界框和标签
cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
label = f"Class {int(cls_id)}: {conf:.2f}"
cv2.putText(img, label, (x1, y1 - 5),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
cv2.imshow("Detection Results", img)
cv2.waitKey(0)
cv2.destroyAllWindows()
if __name__ == "__main__":
# 配置参数
ENGINE_PATH = "yolov5s.engine" # TensorRT引擎文件
IMAGE_PATH = "test_image.jpg" # 测试图像
# 初始化推理器
inferencer = YOLOv5TRTInference(ENGINE_PATH)
# 预处理
input_data = inferencer.preprocess_image(IMAGE_PATH)
# 执行推理并计时
start_time = time.time()
outputs = inferencer.inference(input_data)
inference_time = time.time() - start_time
# 后处理
detections = inferencer.postprocess(outputs)
# 输出结果
print(f"Inference time: {inference_time*1000:.2f} ms")
print(f"Detected {len(detections)} objects:")
for i, det in enumerate(detections):
print(f"Object {i+1}: class={int(det[5])}, conf={det[4]:.4f}, bbox={det[:4]}")
# 可视化结果
inferencer.visualize(IMAGE_PATH, detections)
```
### 关键说明：
1. **环境要求**：
- CUDA 12.6 + TensorRT 10.9
- PyCUDA 2023.1+
- OpenCV 4.7+
- NumPy 1.23+
2. **使用流程**：
```mermaid
graph LR
A[加载TensorRT引擎] --> B[预处理图像]
B --> C[执行GPU推理]
C --> D[后处理/NMS]
D --> E[可视化结果]
```
3. **性能优化**：
- 使用`execute_async_v2`实现异步推理
- 使用`pagelocked memory`加速主机-设备传输
- 支持动态批量大小（需在导出时设置）
4. **注意事项**：
- 确保引擎文件使用相同TensorRT版本导出
- 输入图像尺寸需匹配引擎输入尺寸
- 后处理逻辑需根据实际输出层结构调整
### 引擎导出建议：
使用以下命令导出TensorRT引擎：
```bash
python export.py --weights yolov5s.pt \
--include engine \
--device 0 \
--half \
--opset 18 \
--imgsz 640 \
--iou-thres 0.45 \
--conf-thres 0.5
```
扫一扫
3842
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
