Linux下YOLOv8 TensorRT模型部署

一、背景与问题

在边缘计算和实时推理场景中，模型部署性能是关键指标。YOLOv8作为当前最先进的目标检测算法，其推理速度和精度在多个基准测试中均表现优异。然而，传统PyTorch模型在部署时存在以下几个核心问题：

计算资源占用高：PyTorch模型在运行时需要维护完整的计算图，导致内存占用大
推理速度不足：在CPU上运行时，YOLOv8的FPS通常在10-20帧之间
跨平台兼容性差：在不同硬件平台上需要重新训练和优化模型

TensorRT作为NVIDIA提供的高性能深度学习推理库，通过以下手段解决上述问题：

自动优化计算图（层融合、内存优化）
支持FP16/FP32/INT8精度转换
提供高效的内存管理机制
支持多线程并发处理

在部署YOLOv8模型时，需要完成三个核心步骤：模型转换、引擎构建、推理部署。本文将深入解析这一过程的技术细节。

二、基本原理

1. YOLOv8模型结构

YOLOv8采用改进的CSPDarknet53主干网络，包含：

3个不同尺度的特征提取层（P3/P4/P5）
3个不同尺度的检测头（xyxy, obj, cls）
优化的梯度裁剪机制

模型输入为416x416的图像，输出包含5个检测类别和边界框信息。

2. TensorRT工作原理

TensorRT通过以下技术提升推理性能：

计算图优化：自动合并冗余计算节点，例如将多个矩阵乘法合并为单个操作
内存优化：将模型权重转换为更高效的存储格式（如FP16）
执行计划优化：根据硬件特性选择最优的执行顺序
并发执行：支持多线程处理多个推理请求

三、环境准备

1. 系统要求

Linux系统（Ubuntu 20.04或更高版本）
CUDA 11.8
cuDNN 8.6.0
TensorRT 8.6.3
Python 3.8+（建议3.9）

2. 安装依赖

# 安装系统依赖
sudo apt-get update
sudo apt-get install -y build-essential cmake libgl1 libx11-dev libxext-dev

# 安装CUDA工具包
wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_559.45.00_linux.iso
sudo mount -o loop cuda_11.8.0_559.45.00_linux.iso /mnt
sudo dpkg -i /mnt/*.deb
sudo umount /mnt

# 安装cuDNN
wget https://developer.download.nvidia.com/compute/cudnn/8.6.0/cudnn-linux-x86_64-8.6.0.17.tar.gz
tar -xzf cudnn-linux-x86_64-8.6.0.17.tar.gz
sudo cp cuda/include/cudnn.h /usr/local/include/
sudo cp cuda/lib/libcudnn.so* /usr/local/lib/
sudo ldconfig

# 安装TensorRT
wget https://developer.download.nvidia.com/compute/tensorrt/8.6.3/tensorrt-8.6.3-linux-x64-gnu.tar.gz
tar -xzf tensorrt-8.6.3-linux-x64-gnu.tar.gz
export PATH=$PATH:/usr/local/TensorRT-8.6.3/bin

四、核心实现

1. 模型转换

将YOLOv8 PyTorch模型转换为ONNX格式：

# yolo_to_onnx.py
import torch
from models import YOLOv8

# 加载预训练模型
model = YOLOv8("yolov8m.pt")
model.eval()

# 导出ONNX模型
dummy_input = torch.rand(1, 3, 640, 640)
torch.onnx.export(
    model,
    dummy_input,
    "yolov8m.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
    opset_version=13
)

说明：需要确保models/YOLOv8.py文件存在，该文件包含YOLOv8模型定义。

2. TensorRT引擎构建

# build_trt_engine.py
import tensorrt as trt
import numpy as np
import os

# 加载ONNX模型
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, TRT_LOGGER)

with open("yolov8m.onnx", "rb") as f:
    if not parser.parse(f):
        print("Failed parsing ONNX file")
        for error in parser.get_errors():
            print(error)
        exit()

# 配置构建参数
config = builder.create_builder_config()
config.max_workspace_size = 1 << 30  # 1GB
config.set_flag(trt.BuilderFlag.PERSISTENT_WORKSPACE)

# 构建引擎
engine = builder.build_engine(network, config)
if not engine:
    print("Failed building engine")
    exit()

# 保存引擎文件
with open("yolov8m.trt", "wb") as f:
    f.write(engine.serialize())

3. 推理部署

# inference_trt.py
import tensorrt as trt
import numpy as np
import cv2

# 加载TensorRT引擎
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
runtime = trt.Runtime(TRT_LOGGER)
with open("yolov8m.trt", "rb") as f:
    engine = runtime.deserialize_cuda_engine(f.read())

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

# 创建输入输出张量
input_data = np.empty((1, 3, 640, 640), dtype=np.float16)
output_data = np.empty((1, 84, 84, 85), dtype=np.float16)

# 推理函数
def infer(image):
    # 预处理图像
    img = cv2.resize(image, (640, 640))
    img = img.astype(np.float16) / 255.0
    img = np.transpose(img, (2, 0, 1))  # HWC -> CHW
    
    # 填充输入数据
    np.copyto(input_data, img)
    
    # 执行推理
    context.execute_v2([input_data])
    
    # 获取输出结果
    np.copyto(output_data, output_data)
    
    return output_data

五、完整案例

1. 实时视频检测系统

# yolo_trt_realtime.py
import cv2
import numpy as np
import time

# 加载TensorRT引擎
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
runtime = trt.Runtime(TRT_LOGGER)
with open("yolov8m.trt", "rb") as f:
    engine = runtime.deserialize_cuda_engine(f.read())

context = engine.create_execution_context()

# 加载模型参数
input_data = np.empty((1, 3, 640, 640), dtype=np.float16)
output_data = np.empty((1, 84, 84, 85), dtype=np.float16)

# 打开摄像头
cap = cv2.VideoCapture(0)

while True:
    start_time = time.time()
    
    # 读取帧
    ret, frame = cap.read()
    if not ret:
        break
    
    # 预处理
    img = cv2.resize(frame, (640, 640))
    img = img.astype(np.float16) / 255.0
    img = np.transpose(img, (2, 0, 1))
    
    # 推理
    np.copyto(input_data, img)
    context.execute_v2([input_data])
    np.copyto(output_data, output_data)
    
    # 后处理
    results = post_process(output_data)
    
    # 显示结果
    for result in results:
        x1, y1, x2, y2, confidence, class_id = result
        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(frame, f"{class_id} {confidence:.2f}", (x1, y1-10), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
    
    cv2.imshow("YOLOv8 TensorRT Inference", frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break
    
    # 计算FPS
    print(f"FPS: {1.0 / (time.time() - start_time):.2f}")

cap.release()
cv2.destroyAllWindows()

六、源码解析

1. 模型转换关键点

# ONNX模型解析
if not parser.parse(f):
    print("Failed parsing ONNX file")
    for error in parser.get_errors():
        print(error)
    exit()

trt.OnnxParser用于验证模型的合法性
会检查模型的输入输出维度是否符合预期
需要确保模型的输入格式为CHW（通道在前）

2. 引擎构建关键点

config.max_workspace_size = 1 << 30
config.set_flag(trt.BuilderFlag.PERSISTENT_WORKSPACE)

max_workspace_size控制构建过程中的临时内存大小
PERSISTENT_WORKSPACE标志启用持久化内存优化
实际部署时需根据硬件资源调整此值

3. 推理执行关键点

context.execute_v2([input_data])

execute_v2方法支持多输入/输出的执行
需要确保输入输出的顺序与网络定义一致
使用np.copyto进行数据拷贝时要注意内存对齐

七、进阶使用

1. 精度转换优化

# 构建INT8引擎
config.set_flag(trt.BuilderFlag.INT8)
config.int8_calibrator = MyCalibrator()

使用INT8精度可降低内存占用（约1/4）
需要训练校准数据集（通常为1000张图片）
需要确保输入数据分布与训练数据一致

2. 并发处理优化

# 多线程推理
import threading

def worker():
    while True:
        if not queue.empty():
            frame = queue.get()
            # 执行推理并保存结果

thread = threading.Thread(target=worker)
thread.start()

使用线程池处理多帧数据
可结合cv2.CAP_PROP_FPS控制帧率
需要合理设置线程数量（通常为CPU核心数）

3. 内存管理优化

# 使用内存池管理
memory_pool = trt.MemoryPool(1 << 20)  # 1MB
input_data = memory_pool.allocate(1, 3, 640, 640, np.float16)

减少频繁内存分配/释放的开销
适用于高并发场景
需要确保内存池大小足够

八、性能与工程实践

1. 性能优化策略

优化手段	效果	说明
使用FP16精度	20-30%	减少内存占用，提升计算速度
启用INT8量化	30-50%	降低计算精度，但保持较高精度
启用TensorRT优化	15-20%	自动优化计算图结构
多线程处理	10-20%	提高并发处理能力

2. 异常处理机制

try:
    context.execute_v2([input_data])
except Exception as e:
    print("Inference error:", e)
    # 重置上下文或重新加载引擎

需要处理内存错误、计算图错误等
建议设置超时机制（使用trt.IExecutionContext的setOptimizationProfile）

3. 安全风险控制

模型文件应存储在受保护的目录中
禁用不必要的API接口
对输入数据进行校验和过滤
避免将敏感数据存储在显存中

九、常见问题与踩坑

1. 模型转换失败

错误示例：

parser.parse(f)  # 没有检查返回值

解决方法：

if not parser.parse(f):
    for error in parser.get_errors():
        print(error)

2. 推理速度慢

常见原因：

没有启用TensorRT优化
使用FP32精度
没有正确设置输入输出格式

解决方案：

启用trt.BuilderFlag.FP16优化
确保输入数据为CHW格式
调整max_workspace_size参数

3. 内存不足