1. 项目概述
在计算机视觉领域，YOLO系列算法因其出色的实时检测性能而广受欢迎。而TensorRT作为NVIDIA推出的高性能深度学习推理框架，能够显著提升模型在NVIDIA GPU上的运行效率。本文将详细介绍在Ubuntu 22.04系统上配置TensorRT推理YOLO模型的全套环境，涵盖从基础环境搭建到最终模型部署的完整流程。
这个配置过程涉及多个关键组件：CUDA工具包、cuDNN加速库、TensorRT推理引擎以及OpenCV等视觉处理库的安装与配置。对于需要在实际项目中部署YOLO模型的开发者来说，一个稳定且高效的环境配置是项目成功的基础。
2. 环境准备与依赖安装
2.1 系统基础环境检查
在开始安装前，首先需要确认系统环境是否符合要求：
code复制lsb_release -a # 查看Ubuntu版本
uname -m # 查看系统架构
nvidia-smi # 检查GPU驱动是否安装
确保系统是Ubuntu 22.04 LTS版本，架构为x86_64，并且已安装适当版本的NVIDIA驱动。建议使用470或更高版本的驱动以获得最佳兼容性。
2.2 CUDA工具包安装
TensorRT需要特定版本的CUDA支持。对于Ubuntu 22.04，推荐安装CUDA 11.7版本：
code复制wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-ubuntu2204.pin
sudo mv cuda-ubuntu2204.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/11.7.1/local_installers/cuda-repo-ubuntu2204-11-7-local_11.7.1-515.65.01-1_amd64.deb
sudo dpkg -i cuda-repo-ubuntu2204-11-7-local_11.7.1-515.65.01-1_amd64.deb
sudo cp /var/cuda-repo-ubuntu2204-11-7-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install cuda
安装完成后，将CUDA添加到环境变量中：
code复制echo 'export PATH=/usr/local/cuda-11.7/bin${PATH:+:${PATH}}' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-11.7/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}' >> ~/.bashrc
source ~/.bashrc
2.3 cuDNN安装
cuDNN是NVIDIA提供的深度神经网络加速库，需要与CUDA版本匹配。从NVIDIA官网下载对应版本的cuDNN，然后执行：
code复制sudo dpkg -i libcudnn8_8.x.x.x-1+cuda11.x_amd64.deb
sudo dpkg -i libcudnn8-dev_8.x.x.x-1+cuda11.x_amd64.deb
sudo dpkg -i libcudnn8-samples_8.x.x.x-1+cuda11.x_amd64.deb
安装完成后，可以运行测试样例验证安装是否成功：
code复制cp -r /usr/src/cudnn_samples_v8/ $HOME
cd $HOME/cudnn_samples_v8/mnistCUDNN
make clean && make
./mnistCUDNN
3. TensorRT安装与配置
3.1 TensorRT安装
对于Ubuntu 22.04，推荐安装TensorRT 8.4 GA版本：
code复制sudo apt-get install libnvinfer8 libnvonnxparsers8 libnvparsers8 libnvinfer-plugin8
sudo apt-get install libnvinfer-dev libnvonnxparsers-dev libnvparsers-dev libnvinfer-plugin-dev
sudo apt-get install python3-libnvinfer python3-libnvinfer-dev
也可以从NVIDIA官网下载tar包进行安装，这种方式更加灵活：
code复制tar -xzvf TensorRT-8.x.x.x.Linux.x86_64-gnu.cuda-11.x.cudnn8.x.tar.gz
cd TensorRT-8.x.x.x
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$(pwd)/lib
3.2 验证TensorRT安装
安装完成后，可以通过Python接口验证TensorRT是否正常工作：
python复制import tensorrt as trt
print(trt.__version__)
还可以运行TensorRT自带的样例程序进行更全面的测试：
code复制cd samples/sampleMNIST
make
cd ../../data/mnist
python3 download_pgms.py
cd ../../bin
./sample_mnist
4. YOLO模型转换与优化
4.1 获取YOLO模型
以YOLOv5为例，首先克隆官方仓库并安装依赖：
code复制git clone https://github.com/ultralytics/yolov5.git
cd yolov5
pip install -r requirements.txt
下载预训练模型：
code复制python - <<EOF
from utils.downloads import attempt_download
attempt_download('yolov5s.pt')
EOF
4.2 模型导出为ONNX格式
TensorRT需要通过ONNX格式转换YOLO模型：
code复制python export.py --weights yolov5s.pt --include onnx --imgsz 640 640
导出时需要注意以下几点：
- 指定固定的输入尺寸（如640x640）
- 确保opset_version>=11
- 对于动态batch支持，可以使用--dynamic参数
4.3 ONNX模型优化
使用onnx-simplifier优化导出的ONNX模型：
code复制pip install onnx-simplifier
python -m onnxsim yolov5s.onnx yolov5s-sim.onnx
优化后的模型通常体积更小，推理效率更高。
4.4 转换为TensorRT引擎
使用TensorRT的trtexec工具将ONNX模型转换为TensorRT引擎：
code复制/usr/src/tensorrt/bin/trtexec --onnx=yolov5s-sim.onnx --saveEngine=yolov5s.engine --fp16
关键参数说明：
- --fp16：启用FP16精度，可显著提升推理速度
- --workspace：设置最大工作空间大小（默认为16MB，对于YOLO建议设置为1-2GB）
- --minShapes/--optShapes/--maxShapes：设置动态形状的尺寸范围
5. 推理代码实现
5.1 Python推理接口
创建一个基本的YOLO推理类：
python复制import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import cv2
class YOLO_TRT:
def __init__(self, engine_path):
self.logger = trt.Logger(trt.Logger.WARNING)
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 = 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)
self.bindings.append(int(device_mem))
if self.engine.binding_is_input(binding):
self.inputs.append({'host': host_mem, 'device': device_mem})
else:
self.outputs.append({'host': host_mem, 'device': device_mem})
def infer(self, img):
# 预处理
img = cv2.resize(img, (640, 640))
img = img.transpose((2, 0, 1)) # HWC to CHW
img = np.ascontiguousarray(img)
img = img.astype(np.float32) / 255.0
# 拷贝数据到GPU
np.copyto(self.inputs[0]['host'], img.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)
# 拷贝结果回CPU
cuda.memcpy_dtoh_async(self.outputs[0]['host'], self.outputs[0]['device'], self.stream)
self.stream.synchronize()
# 后处理
output = self.outputs[0]['host']
return output
5.2 后处理实现
YOLO的输出需要经过非极大值抑制(NMS)等后处理：
python复制def postprocess(output, conf_thres=0.25, iou_thres=0.45):
# output shape: (1, 25200, 85)
output = output.reshape(1, -1, 85)
# 过滤低置信度检测
mask = output[..., 4] > conf_thres
output = output[mask]
# 计算类别分数
scores = output[..., 5:] * output[..., 4:5]
# 转换为xywh格式
boxes = output[..., :4]
boxes[..., 0] -= boxes[..., 2] / 2 # x center to xmin
boxes[..., 1] -= boxes[..., 3] / 2 # y center to ymin
# NMS处理
keep = []
for cls in range(scores.shape[1]):
mask = scores[:, cls] > conf_thres
if not mask.any():
continue
cls_boxes = boxes[mask]
cls_scores = scores[mask, cls]
indices = cv2.dnn.NMSBoxes(cls_boxes.tolist(), cls_scores.tolist(),
conf_thres, iou_thres)
keep.extend([(i, cls) for i in indices.flatten()])
return boxes[keep], scores[keep]
6. 性能优化技巧
6.1 精度选择与量化
TensorRT支持多种精度模式：
- FP32：最高精度，速度最慢
- FP16：平衡精度和速度，推荐大多数场景使用
- INT8：最高速度，需要校准
启用FP16模式可以显著提升性能：
code复制builder.fp16_mode = True
对于INT8量化，需要准备校准数据集：
python复制class Calibrator(trt.IInt8EntropyCalibrator2):
def __init__(self, calibration_data):
super().__init__()
self.data = calibration_data
self.current_index = 0
def get_batch_size(self):
return 1
def get_batch(self, names):
if self.current_index >= len(self.data):
return None
batch = self.data[self.current_index]
self.current_index += 1
return [batch.data_ptr()]
6.2 动态形状支持
对于需要处理不同输入尺寸的场景，可以启用动态形状：
python复制profile = builder.create_optimization_profile()
profile.set_shape("input", (1,3,320,320), (1,3,640,640), (1,3,1280,1280))
config.add_optimization_profile(profile)
6.3 层融合与图优化
TensorRT会自动执行以下优化：
- 垂直融合：将多个层合并为一个复合层
- 水平融合：将并行执行的层合并
- 消除无用操作：如恒等变换、零操作等
可以通过以下方式查看优化后的网络：
python复制for i in range(engine.num_bindings):
name = engine.get_binding_name(i)
dtype = engine.get_binding_dtype(i)
shape = engine.get_binding_shape(i)
print(f"Binding {i}: name={name}, dtype={dtype}, shape={shape}")
7. 常见问题与解决方案
7.1 CUDA版本不兼容
错误现象：
code复制Could not load library libcudnn.so.8
解决方案：
- 检查CUDA和cuDNN版本是否匹配
- 确保LD_LIBRARY_PATH包含CUDA和cuDNN库路径
- 使用
ldconfig -p | grep cudnn验证库是否被系统识别
7.2 TensorRT推理结果异常
可能原因：
- ONNX导出时参数设置不正确
- 输入数据预处理不一致
- 输出后处理逻辑错误
调试步骤：
- 使用ONNX Runtime验证ONNX模型输出
- 比较TensorRT和PyTorch原始模型的输出差异
- 检查输入数据的归一化和通道顺序
7.3 性能未达预期
优化建议：
- 使用
nvprof分析瓶颈 - 尝试不同的精度模式（FP32/FP16/INT8）
- 调整工作空间大小
- 启用TF32（Ampere架构及以上GPU）
7.4 内存不足问题
处理方法：
- 减少batch size
- 使用更小的模型尺寸
- 启用内存池：
python复制config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30) # 1GB
8. 完整部署示例
8.1 实时摄像头检测
python复制import time
def main():
# 初始化模型
yolo = YOLO_TRT("yolov5s.engine")
# 打开摄像头
cap = cv2.VideoCapture(0)
while True:
start = time.time()
ret, frame = cap.read()
if not ret:
break
# 推理
output = yolo.infer(frame)
boxes, scores = postprocess(output)
# 绘制结果
for box, score in zip(boxes, scores):
x1, y1, w, h = box
cv2.rectangle(frame, (int(x1), int(y1)),
(int(x1+w), int(y1+h)), (0,255,0), 2)
# 显示FPS
fps = 1 / (time.time() - start)
cv2.putText(frame, f"FPS: {fps:.2f}", (10,30),
cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,255), 2)
cv2.imshow("YOLO TensorRT", frame)
if cv2.waitKey(1) == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
if __name__ == "__main__":
main()
8.2 批量图像处理
对于需要处理大量静态图像的应用：
python复制from pathlib import Path
def batch_inference(image_dir, output_dir):
yolo = YOLO_TRT("yolov5s.engine")
image_paths = list(Path(image_dir).glob("*.jpg"))
for img_path in image_paths:
img = cv2.imread(str(img_path))
output = yolo.infer(img)
boxes, scores = postprocess(output)
# 绘制并保存结果
for box, score in zip(boxes, scores):
x1, y1, w, h = box
cv2.rectangle(img, (int(x1), int(y1)),
(int(x1+w), int(y1+h)), (0,255,0), 2)
output_path = Path(output_dir) / img_path.name
cv2.imwrite(str(output_path), img)
9. 进阶配置与优化
9.1 多流并行处理
利用CUDA流实现并行推理：
python复制class YOLO_TRT_MultiStream:
def __init__(self, engine_path, num_streams=2):
self.logger = trt.Logger(trt.Logger.WARNING)
with open(engine_path, "rb") as f, trt.Runtime(self.logger) as runtime:
engine_data = f.read()
self.engines = []
self.contexts = []
self.streams = []
for _ in range(num_streams):
engine = runtime.deserialize_cuda_engine(engine_data)
context = engine.create_execution_context()
stream = cuda.Stream()
self.engines.append(engine)
self.contexts.append(context)
self.streams.append(stream)
# 为每个流分配内存
self.bindings_list = []
for engine in self.engines:
bindings = []
for binding in engine:
size = trt.volume(engine.get_binding_shape(binding))
dtype = trt.nptype(engine.get_binding_dtype(binding))
host_mem = cuda.pagelocked_empty(size, dtype)
device_mem = cuda.mem_alloc(host_mem.nbytes)
bindings.append({'host': host_mem, 'device': device_mem})
self.bindings_list.append(bindings)
9.2 自定义插件支持
对于YOLO中的特殊层（如SiLU激活函数），可能需要注册自定义插件：
python复制class SiLUPlugin(trt.IPluginV2DynamicExt):
def __init__(self):
super().__init__()
def get_output_datatype(self, index, input_types):
return input_types[0]
def configure_plugin(self, in_pos, in_types, in_shapes, out_pos, out_types, out_shapes):
pass
def initialize(self):
return 0
def terminate(self):
pass
def get_serialization_size(self):
return 0
def serialize(self):
return b""
def destroy(self):
pass
# 注册插件
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
trt.init_libnvinfer_plugins(TRT_LOGGER, "")
registry = trt.get_plugin_registry()
registry.register_creator(SiLUPlugin, "", "")
9.3 性能监控与分析
使用NVIDIA Nsight Systems进行性能分析：
code复制nsys profile -o yolov5_profile python yolov5_trt.py
分析结果可以显示：
- GPU利用率
- 内存拷贝时间
- 核函数执行时间
- 各层计算耗时
10. 环境维护与更新
10.1 版本兼容性矩阵
保持各组件版本兼容至关重要：
| 组件 | 推荐版本 | 兼容范围 |
|---|---|---|
| Ubuntu | 22.04 LTS | 20.04-22.04 |
| NVIDIA驱动 | 515.65.01 | >=470 |
| CUDA | 11.7 | 11.4-11.8 |
| cuDNN | 8.5.0 | 8.3.0-8.6.0 |
| TensorRT | 8.4.3 | 8.2.0-8.5.0 |
| PyTorch | 1.12.1 | 1.10.0-1.13.0 |
10.2 容器化部署
使用NVIDIA官方容器简化环境配置：
code复制docker pull nvcr.io/nvidia/tensorrt:22.07-py3
docker run --gpus all -it --rm -v $(pwd):/workspace nvcr.io/nvidia/tensorrt:22.07-py3
容器内已预装：
- CUDA 11.7
- cuDNN 8.5.0
- TensorRT 8.4.3
- PyTorch 1.12.0
10.3 自动化部署脚本
创建一键部署脚本setup.sh：
bash复制#!/bin/bash
# 安装基础依赖
sudo apt-get update
sudo apt-get install -y build-essential cmake git wget
# 安装CUDA
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-ubuntu2204.pin
sudo mv cuda-ubuntu2204.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/11.7.1/local_installers/cuda-repo-ubuntu2204-11-7-local_11.7.1-515.65.01-1_amd64.deb
sudo dpkg -i cuda-repo-ubuntu2204-11-7-local_11.7.1-515.65.01-1_amd64.deb
sudo cp /var/cuda-repo-ubuntu2204-11-7-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install cuda
# 安装cuDNN
sudo apt-get install -y libcudnn8 libcudnn8-dev libcudnn8-samples
# 安装TensorRT
sudo apt-get install -y libnvinfer8 libnvonnxparsers8 libnvparsers8 libnvinfer-plugin8 \
libnvinfer-dev libnvonnxparsers-dev libnvparsers-dev libnvinfer-plugin-dev \
python3-libnvinfer python3-libnvinfer-dev
# 配置环境变量
echo 'export PATH=/usr/local/cuda-11.7/bin${PATH:+:${PATH}}' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-11.7/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}' >> ~/.bashrc
source ~/.bashrc
# 验证安装
nvcc --version
python3 -c "import tensorrt; print(tensorrt.__version__)"