Ubuntu 22.04配置TensorRT实现YOLO模型高效推理

1. 环境配置背景与核心需求

在计算机视觉领域，YOLO系列算法因其出色的实时检测性能而广受欢迎。而TensorRT作为NVIDIA推出的高性能推理引擎，能够显著提升模型在NVIDIA GPU上的运行效率。Ubuntu 22.04 LTS作为当前最稳定的Linux发行版之一，为深度学习开发提供了良好的基础环境。

这个配置方案主要解决三个核心问题：

1. 如何在Ubuntu 22.04上搭建完整的TensorRT推理环境
2. 如何将YOLO模型转换为TensorRT引擎并优化推理性能
3. 如何验证环境配置的正确性和推理效果

注意：建议使用NVIDIA 30系或更新架构的GPU，因为TensorRT对新一代GPU的支持更完善。如果使用20系或更早的GPU，部分功能可能受限。

2. 基础环境准备

2.1 系统与驱动安装

首先确保系统为Ubuntu 22.04.3 LTS或更新版本。安装完成后，需要配置NVIDIA驱动：

bash复制# 查看推荐驱动版本
ubuntu-drivers devices

# 安装推荐驱动（示例为525版本）
sudo apt install nvidia-driver-525

安装完成后重启系统，通过nvidia-smi命令验证驱动安装是否成功。正常情况应显示GPU信息和驱动版本。

2.2 CUDA Toolkit安装

TensorRT 8.x需要CUDA 11.x支持，而TensorRT 10.x需要CUDA 12.x。这里以TensorRT 8.6 + CUDA 11.8为例：

bash复制wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run
sudo sh cuda_11.8.0_520.61.05_linux.run

安装时注意：

• 不要安装驱动（已单独安装）
• 确保安装CUDA Samples（用于验证）
• 添加环境变量到~/.bashrc：

bash复制export PATH=/usr/local/cuda-11.8/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}

2.3 cuDNN安装

下载对应CUDA 11.8的cuDNN 8.9.x版本，执行：

bash复制sudo dpkg -i libcudnn8_8.9.4.25-1+cuda11.8_amd64.deb
sudo dpkg -i libcudnn8-dev_8.9.4.25-1+cuda11.8_amd64.deb

3. TensorRT安装与配置

3.1 TensorRT本地安装

下载TensorRT 8.6 GA for Ubuntu 22.04的.deb包，执行：

bash复制sudo dpkg -i nv-tensorrt-local-repo-ubuntu2204-8.6.1-cuda-11.8_1.0-1_amd64.deb
sudo cp /var/nv-tensorrt-local-repo-ubuntu2204-8.6.1-cuda-11.8/trt-8.6.1.6-ga-20231005.gpg /usr/share/keyrings/
sudo apt update
sudo apt install tensorrt

验证安装：

bash复制dpkg -l | grep TensorRT

3.2 Python接口安装

推荐使用Python 3.10虚拟环境：

bash复制python3.10 -m venv trt_env
source trt_env/bin/activate
pip install tensorrt-8.6.1-cp310-none-linux_x86_64.whl
pip install pycuda

提示：如果使用TensorRT 10.x，需要对应Python 3.11环境

4. YOLO模型转换与优化

4.1 ONNX模型导出

以YOLOv8为例，首先导出ONNX模型：

python复制from ultralytics import YOLO

model = YOLO('yolov8n.pt')  # 加载预训练模型
model.export(format='onnx', dynamic=True, simplify=True)  # 导出动态尺寸ONNX

关键参数说明：

• dynamic=True：允许可变输入尺寸
• simplify=True：简化模型结构
• opset=13：建议使用ONNX opset 13或更高

4.2 TensorRT引擎生成

使用trtexec工具转换：

bash复制trtexec --onnx=yolov8n.onnx \
        --saveEngine=yolov8n.engine \
        --fp16 \
        --workspace=4096 \
        --verbose

优化参数解析：

• --fp16：启用FP16精度，提升推理速度
• --workspace：设置显存工作区大小(MB)
• --best：自动选择最优策略（TensorRT 8.6+）

4.3 动态尺寸处理

对于需要支持多尺寸输入的情况，需明确指定优化范围：

bash复制trtexec --onnx=yolov8n.onnx \
        --minShapes=images:1x3x320x320 \
        --optShapes=images:1x3x640x640 \
        --maxShapes=images:1x3x1280x1280 \
        --saveEngine=yolov8n_dynamic.engine

5. 推理代码实现

5.1 Python推理示例

python复制import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np

class YOLOTRT:
    def __init__(self, engine_path):
        self.logger = trt.Logger(trt.Logger.WARNING)
        with open(engine_path, "rb") as f:
            self.engine = trt.Runtime(self.logger).deserialize_cuda_engine(f.read())
        self.context = self.engine.create_execution_context()
        
        # 获取输入输出绑定
        self.bindings = []
        for binding in self.engine:
            size = trt.volume(self.engine.get_binding_shape(binding)) 
            dtype = trt.nptype(self.engine.get_binding_dtype(binding))
            self.bindings.append(cuda.mem_alloc(size * dtype.itemsize))
            
    def infer(self, input_data):
        # 数据传输到GPU
        cuda.memcpy_htod(self.bindings[0], input_data)
        
        # 执行推理
        self.context.execute_v2(bindings=self.bindings)
        
        # 获取输出
        outputs = []
        for i in range(1, self.engine.num_bindings):
            output = np.empty(self.engine.get_binding_shape(i), 
                            dtype=trt.nptype(self.engine.get_binding_dtype(i)))
            cuda.memcpy_dtoh(output, self.bindings[i])
            outputs.append(output)
        return outputs

5.2 后处理实现

YOLO输出需要NMS等后处理：

python复制def postprocess(outputs, conf_thres=0.5, iou_thres=0.5):
    # 解析TensorRT输出
    predictions = outputs[0].reshape(-1, 85)  # 假设输出格式为[x,y,w,h,conf,cls1,cls2...]
    
    # 过滤低置信度检测
    mask = predictions[:, 4] > conf_thres
    predictions = predictions[mask]
    
    # 执行NMS
    boxes = predictions[:, :4]
    scores = predictions[:, 4] * predictions[:, 5:].max(1)
    keep = nms(boxes, scores, iou_thres)
    
    return predictions[keep]

6. 性能优化技巧

6.1 精度选择策略

建议开发阶段使用FP16，部署时根据需求考虑INT8量化：

bash复制trtexec --onnx=yolov8n.onnx \
        --saveEngine=yolov8n_int8.engine \
        --int8 \
        --calib=calibration.cache

6.2 多流并行处理

python复制# 创建多个执行上下文
contexts = [engine.create_execution_context() for _ in range(4)]

# 使用不同流处理
streams = [cuda.Stream() for _ in range(4)]
for i, (input_data, context, stream) in enumerate(zip(input_batch, contexts, streams)):
    cuda.memcpy_htod_async(bindings[0], input_data, stream)
    context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
    cuda.memcpy_dtoh_async(outputs[i], bindings[1], stream)

7. 常见问题排查

7.1 版本兼容性问题

常见错误组合：

• TensorRT 8.x + CUDA 11.x + cuDNN 8.x
• TensorRT 10.x + CUDA 12.x + cuDNN 9.x

验证命令：

bash复制nvcc --version  # CUDA版本
cat /usr/include/cudnn_version.h | grep CUDNN_MAJOR -A 2  # cuDNN版本
python -c "import tensorrt; print(tensorrt.__version__)"  # TensorRT版本

7.2 模型转换失败处理

典型错误及解决方案：

7.3 推理精度异常

排查步骤：

1. 对比ONNX和TensorRT引擎输出差异
2. 检查FP16/INT8量化是否导致精度损失过大
3. 验证输入数据预处理是否一致

python复制# 对比原始PyTorch和TensorRT输出
torch_output = torch_model(torch_input).detach().numpy()
trt_output = trt_model(np_input)

print("最大差异:", np.max(np.abs(torch_output - trt_output)))

8. 环境验证与基准测试

8.1 基础功能验证

使用TensorRT自带的sample_mnist测试：

bash复制cd /usr/src/tensorrt/samples/sample_mnist
make
./sample_mnist

8.2 YOLO推理基准

测试脚本示例：

python复制import time

# 预热
for _ in range(10):
    trt_model.infer(test_input)

# 正式测试
start = time.time()
for _ in range(100):
    trt_model.infer(test_input)
elapsed = (time.time() - start)/100

print(f"平均推理时间: {elapsed*1000:.2f}ms")
print(f"FPS: {1/elapsed:.2f}")

典型性能参考（RTX 3090, YOLOv8n）：

8.3 显存监控

使用nvidia-smi -l 1观察显存占用变化，确保没有内存泄漏。正常情况显存应在推理时稳定波动，不会持续增长。