使用DAMO-YOLO和TensorRT实现高性能推理加速-深圳市維司達科技有限公司

使用DAMO-YOLO和TensorRT实现高性能推理加速

最近在做一个实时视频分析的项目，对目标检测模型的推理速度要求特别高。用PyTorch直接跑DAMO-YOLO，虽然精度不错，但帧率总上不去，离真正的“实时”还差点意思。相信很多做边缘部署或者对延迟敏感的朋友都遇到过类似问题。

后来我把目光投向了TensorRT，这是NVIDIA官方推出的高性能推理优化器。经过一番折腾，成功把DAMO-YOLO的推理速度提升了好几倍，而且精度损失微乎其微。整个过程有点像给模型做“瘦身”和“加速”，让它能在资源有限的设备上跑得更快更稳。

今天这篇文章，我就把自己从模型转换、精度校准到最终部署的完整流程梳理一遍。如果你也在为模型推理速度发愁，或者想把手头的检测模型部署到Jetson这类边缘设备上，那这篇实战指南应该能帮到你。咱们不聊太多理论，直接上手操作。

1. 环境准备与工具安装

工欲善其事，必先利其器。在开始优化之前，我们需要把必要的软件环境搭建好。整个过程主要围绕Python和NVIDIA的一系列工具展开。

1.1 基础环境检查

首先，确保你有一张支持CUDA的NVIDIA显卡。打开终端，用下面这个命令看看你的显卡信息：

nvidia-smi

你会看到一个表格，里面显示了显卡型号、驱动版本以及CUDA版本。记住你的CUDA版本号，比如11.7或12.1，这决定了后续要安装的TensorRT版本。

接下来，我建议创建一个独立的Python虚拟环境，避免和系统里其他项目的包版本冲突。用conda或者venv都可以，这里以conda为例：

# 创建一个名为 damo_trt 的新环境，指定Python版本 conda create -n damo_trt python=3.8 -y conda activate damo_trt

1.2 核心工具安装

环境激活后，开始安装核心的PyTorch和TensorRT。这里要特别注意版本匹配，尤其是PyTorch的CUDA版本需要和你系统安装的CUDA版本一致。

假设你的CUDA是11.7，可以这样安装PyTorch：

pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117

安装完PyTorch，接下来是重头戏TensorRT。最省事的方法是去NVIDIA官网下载对应你系统和CUDA版本的TensorRT本地安装包（.tar.gz格式）。解压后，将其中的lib文件夹路径添加到环境变量，并把python文件夹下的whl包安装到你的虚拟环境里。

# 假设你把TensorRT解压到了 /home/user/TensorRT-8.5.3.1 export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/user/TensorRT-8.5.3.1/lib # 进入python目录，安装对应的whl包 cd /home/user/TensorRT-8.5.3.1/python pip install tensorrt-8.5.3.1-cp38-none-linux_x86_64.whl

此外，我们还需要两个辅助工具：pycuda用于底层CUDA操作，onnx用于模型格式转换。

pip install pycuda onnx

最后，把DAMO-YOLO的官方代码库克隆下来，并安装其依赖。

git clone https://github.com/tinyvision/DAMO-YOLO.git cd DAMO-YOLO pip install -r requirements.txt

至此，所有准备工作就完成了。你可以运行一个简单的import torch和import tensorrt来测试环境是否正常。

2. 从PyTorch到ONNX：模型转换第一步

TensorRT不能直接读取PyTorch的.pth模型文件，它需要一个中间格式。ONNX（Open Neural Network Exchange）就是这个理想的“桥梁”。这一步的目标是把训练好的DAMO-YOLO模型转换成ONNX格式，同时确保转换过程没有出错。

2.1 准备PyTorch模型

首先，你需要一个训练好的DAMO-YOLO模型权重文件（比如damoyolo_tinynasL20_T.pth）。把它放在一个方便引用的位置。然后，我们需要写一个脚本，按照DAMO-YOLO官方的方式加载这个模型。

关键点在于，PyTorch模型在训练（model.train()）和推理（model.eval()）时的行为是不同的。为了转换，我们必须将模型设置为推理模式。此外，很多现代检测模型（包括DAMO-YOLO）都使用了动态尺寸或后处理技巧，在导出ONNX时需要特别注意。

下面是一个基础的导出脚本框架：

import torch from damo_yolo.models import build_model from damo_yolo.config import get_config # 1. 加载模型配置 config_path = './configs/damoyolo_tinynasL20_T.py' cfg = get_config(config_path, use_amp=False) # 2. 构建模型结构 model = build_model(cfg, test_size=cfg.test_size) # 3. 加载训练好的权重 checkpoint = torch.load('damoyolo_tinynasL20_T.pth', map_location='cpu') model.load_state_dict(checkpoint['model'], strict=True) # 4. 切换到推理模式！这很重要。 model.eval() # 5. 创建一个示例输入张量（模拟一张图片） # 这里的尺寸 [1, 3, 640, 640] 代表：批大小1，3通道RGB，高640，宽640 example_input = torch.randn(1, 3, 640, 640, device='cpu') # 6. 执行一次前向传播，确保模型能跑通 with torch.no_grad(): output = model(example_input) print("PyTorch模型推理成功，输出形状：", output[0].shape if isinstance(output, tuple) else output.shape)

运行这个脚本，如果没有报错，就说明你的PyTorch模型加载正确，可以进入下一步了。

2.2 执行ONNX导出

现在，我们使用PyTorch自带的torch.onnx.export函数进行转换。这里有几个参数需要仔细设置：

opset_version: ONNX算子集版本，建议用12或以上，兼容性更好。
input_names和output_names: 给输入输出起个名字，方便后续识别。
dynamic_axes: 如果你想模型支持动态的输入尺寸（比如高度宽度可变），就在这里指定。但为了初次转换简单，我们先固定尺寸。

import torch.onnx # 指定导出的ONNX文件名 onnx_model_path = "damoyolo_tinynasL20_T.onnx" # 执行导出 torch.onnx.export( model, # 要转换的模型 example_input, # 示例输入 onnx_model_path, # 输出文件路径 opset_version=12, # ONNX算子集版本 input_names=['images'], # 输入节点名 output_names=['outputs'], # 输出节点名 dynamic_axes=None # 首次导出，我们先不用动态尺寸 ) print(f"模型已成功导出至: {onnx_model_path}")

导出完成后，强烈建议用ONNX官方工具onnxruntime或者onnx包里的检查器验证一下文件是否有效。

import onnx # 加载并检查模型 onnx_model = onnx.load(onnx_model_path) onnx.checker.check_model(onnx_model) print("ONNX模型检查通过，格式有效！") # 你也可以简单打印模型信息 print(f"模型输入: {[input.name for input in onnx_model.graph.input]}") print(f"模型输出: {[output.name for output in onnx_model.graph.output]}")

如果一切顺利，你就得到了一个标准的damoyolo_tinynasL20_T.onnx文件。这是通往TensorRT高速推理的“门票”。

3. 使用TensorRT Builder构建优化引擎

拿到ONNX模型后，真正的加速魔法就要开始了。TensorRT的核心是一个叫做Builder的组件，它会分析你的网络结构，进行层融合、精度校准、内核自动选择等一系列优化，最终生成一个高度优化的推理引擎（.engine文件）。这个过程我们称之为“构建”（Build）。

3.1 创建TensorRT Builder和网络

首先，我们需要导入TensorRT库，并创建构建器、网络定义等核心对象。

import tensorrt as trt # 0. 创建一个日志记录器（Logger），用于捕获构建过程中的信息 TRT_LOGGER = trt.Logger(trt.Logger.WARNING) # 1. 创建构建器（Builder），这是优化的总指挥 builder = trt.Builder(TRT_LOGGER) # 2. 创建网络定义（Network）。这里我们显式指定使用最新的“显式Batch”模式。 # 显式Batch模式性能更好，也是现在推荐的方式。 network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) # 3. 创建一个ONNX解析器（Parser），用来读取我们上一步导出的.onnx文件 parser = trt.OnnxParser(network, TRT_LOGGER) onnx_file_path = 'damoyolo_tinynasL20_T.onnx' # 4. 解析ONNX文件 with open(onnx_file_path, 'rb') as model: if not parser.parse(model.read()): print('ONNX解析失败！') for error in range(parser.num_errors): print(parser.get_error(error)) exit() print('ONNX模型解析成功！')

3.2 配置构建参数

接下来，我们需要创建一个构建配置（BuilderConfig），在这里面设置各种优化选项。这是影响最终引擎性能和精度的关键一步。

# 1. 创建构建配置 config = builder.create_builder_config() # 2. 设置最大工作空间大小（单位：字节）。 # 这是TensorRT进行层优化时可以使用的临时GPU内存。设大一点通常有助于找到更优的kernel，但别超过显卡可用内存。 config.max_workspace_size = 1 << 30 # 1 GB # 3. 设置优化级别。 # TensorRT提供了从0到5的优化级别，级别越高，优化时间越长，但可能生成更快的引擎。 # 对于部署，我们通常选择默认或较高级别。 builder.optimization_level = 3 # 4. （可选但重要）精度设置。 # TensorRT支持FP32（单精度）、FP16（半精度）和INT8（整型8位）推理。 # FP16能在几乎不损失精度的情况下大幅提升速度，特别适合现代GPU（如Turing/Ampere架构）。 # 如果你的显卡支持FP16（基本上近几年的都支持），强烈建议开启。 if builder.platform_has_fast_fp16: config.set_flag(trt.BuilderFlag.FP16) print("已启用FP16精度模式。")

关于INT8精度，它能带来更大的加速比，但需要提供一个“校准数据集”来计算每一层激活值的动态范围，过程稍复杂。我们可以在基础教程之后再去探索。对于DAMO-YOLO，FP16通常已经能在速度和精度间取得很好的平衡。

3.3 构建并保存引擎

配置好后，就可以启动构建过程了。Builder会遍历网络，尝试所有可能的优化策略，这可能需要几十秒到几分钟。

# 1. 构建引擎 # 这一步会根据网络定义和配置，生成最优的推理引擎。 # `builder.build_engine(network, config)` 这个API在较新版本中已变更。 # 更通用的方法是使用 `builder.build_serialized_network` serialized_engine = builder.build_serialized_network(network, config) if serialized_engine is None: print("引擎构建失败！") exit() # 2. 将构建好的引擎序列化，保存到.engine文件 engine_file_path = 'damoyolo_tinynasL20_T.engine' with open(engine_file_path, 'wb') as f: f.write(serialized_engine) print(f"TensorRT引擎构建成功，已保存至: {engine_file_path}") print(f"输入绑定名称: {network.get_input(0).name}") print(f"输出绑定名称: {network.get_output(0).name}")

保存下来的.engine文件是特定于你当前GPU架构和TensorRT版本的，不能直接跨平台或跨显卡型号使用。但它也是最终部署时加载速度最快的格式。

4. 加载引擎与执行推理

引擎构建好之后，使用起来就非常高效了。我们需要创建一个Runtime来加载它，然后准备输入输出数据，最后在GPU上执行推理。

4.1 加载引擎并创建执行上下文

import tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit # 初始化CUDA上下文 import numpy as np # 1. 创建Runtime runtime = trt.Runtime(TRT_LOGGER) # 2. 从文件反序列化加载引擎 with open('damoyolo_tinynasL20_T.engine', 'rb') as f: serialized_engine = f.read() engine = runtime.deserialize_cuda_engine(serialized_engine) # 3. 创建执行上下文（Context） # 上下文是实际执行推理的地方。一个引擎可以创建多个上下文（用于多流推理）。 context = engine.create_execution_context() # 4. 获取输入输出的绑定信息（名字、尺寸、数据类型） # 我们需要知道模型期望的输入大小，以便准备数据。 input_binding_name = engine.get_binding_name(0) # 通常第一个是输入 output_binding_name = engine.get_binding_name(1) # 第二个是输出 # 获取输入输出的维度信息 input_shape = engine.get_binding_shape(input_binding_name) # 例如 (1, 3, 640, 640) output_shape = engine.get_binding_shape(output_binding_name) # 输出的形状 print(f"输入绑定: {input_binding_name}, 形状: {input_shape}") print(f"输出绑定: {output_binding_name}, 形状: {output_shape}")

4.2 准备输入输出缓冲区

TensorRT推理需要在GPU内存中分配缓冲区。我们需要在CPU端（Host）准备数据，然后拷贝到GPU端（Device），推理完成后再把结果拷回CPU。

# 1. 根据输入形状，在CPU上准备一个假的输入数据（numpy数组） # 实际应用中，这里应该是你预处理好的图片数据。 batch_size, channel, height, width = input_shape host_input = np.random.random((batch_size, channel, height, width)).astype(np.float32) # 2. 在GPU上分配输入输出内存 # 首先计算需要多少字节 input_nbytes = host_input.nbytes # 假设我们不知道输出具体大小，可以先用引擎信息计算最大可能尺寸，或者先分配一个预估的 # 更规范的做法是根据`output_shape`来分配。对于DAMO-YOLO，输出通常是固定大小的张量。 host_output = np.empty(output_shape, dtype=np.float32) # 在CPU上分配输出内存 output_nbytes = host_output.nbytes # 使用pycuda在GPU上分配内存 device_input = cuda.mem_alloc(input_nbytes) device_output = cuda.mem_alloc(output_nbytes) # 3. 创建CUDA流，用于管理异步的内存拷贝和内核执行 stream = cuda.Stream()

4.3 执行推理并获取结果

一切准备就绪，现在可以运行模型了。

# 1. 将CPU输入数据异步拷贝到GPU cuda.memcpy_htod_async(device_input, host_input, stream) # 2. 执行推理 # 我们需要将输入输出缓冲区的设备指针绑定到上下文的对应索引上 bindings = [int(device_input), int(device_output)] context.execute_async_v2(bindings=bindings, stream_handle=stream.handle) # 3. 将GPU上的推理结果异步拷贝回CPU cuda.memcpy_dtoh_async(host_output, device_output, stream) # 4. 等待流中的所有操作完成 stream.synchronize() print("推理完成！") print(f"输出数据的形状: {host_output.shape}") print(f"输出数据示例（前10个元素）: {host_output.flatten()[:10]}")

4.4 与原始PyTorch结果对比（可选但推荐）

为了确保TensorRT优化后的模型输出是正确的，我们可以将同一个输入分别送入PyTorch模型和TensorRT引擎，然后比较它们的输出是否接近。

# 使用之前加载的PyTorch模型和相同的输入 with torch.no_grad(): torch_output = model(torch.from_numpy(host_input).float()) # 将PyTorch输出转为numpy，并调整形状以便与TensorRT输出比较 # 注意：由于模型后处理或输出格式可能不同，这里可能需要一些reshape或切片操作 # 假设我们只比较第一个输出，且形状一致 torch_output_np = torch_output[0].numpy() if isinstance(torch_output, tuple) else torch_output.numpy() # 计算差异 diff = np.abs(torch_output_np - host_output).max() print(f"PyTorch与TensorRT输出最大绝对误差: {diff}") # 通常，由于FP16精度或优化带来的微小数值差异是允许的。 # 如果误差在1e-3量级或以下，通常可以认为转换是成功的。 if diff < 1e-3: print("✓ 精度验证通过，转换成功！") else: print(" 输出差异较大，请检查模型转换或精度设置。")

5. 性能测试与实用技巧

引擎跑起来之后，我们最关心的就是：到底快了多少？这里有一些测试方法和实用建议。

5.1 如何进行可靠的性能测试

不要只测一次，因为第一次运行可能包含GPU初始化、缓存加载等开销。一个常见的做法是预热（Warm-up）然后多次测量取平均。

import time # 预热 for _ in range(10): context.execute_async_v2(bindings=bindings, stream_handle=stream.handle) stream.synchronize() # 正式计时 num_runs = 100 start_time = time.perf_counter() for _ in range(num_runs): context.execute_async_v2(bindings=bindings, stream_handle=stream.handle) stream.synchronize() end_time = time.perf_counter() # 计算平均耗时和帧率 avg_latency_ms = (end_time - start_time) * 1000 / num_runs fps = 1000 / avg_latency_ms print(f"平均推理延迟: {avg_latency_ms:.2f} ms") print(f"等效帧率 (FPS): {fps:.2f}")

你可以用同样的方法测试原始PyTorch模型的推理速度，然后对比一下。在我的测试环境（RTX 3080）下，DAMO-YOLO-Tiny模型使用TensorRT FP16后，速度相比PyTorch FP32提升了大约2-3倍。

5.2 几个提升体验的实用技巧

动态形状支持：如果你需要处理不同尺寸的图片，可以在构建引擎时通过config.profile设置动态尺寸范围，并在执行推理前通过context.set_binding_shape来指定本次推理的具体尺寸。
使用TensorRT的Python API进行预处理：你可以尝试使用TensorRT的IIdentityLayer或自定义插件，将图像归一化、BGR2RGB等预处理操作也放到GPU上，进一步减少CPU到GPU的数据传输和整体延迟。
多流推理：对于需要同时处理多个视频流的应用，可以为每个CUDA流创建一个执行上下文，实现并行推理，充分利用GPU。
引擎缓存：.engine文件的加载虽然比构建快，但仍有开销。对于长期运行的服务，应该在启动时加载一次引擎，然后重复使用。
关注内存：在Jetson等边缘设备上，内存非常宝贵。构建引擎时，适当调低max_workspace_size，并在运行时监控内存使用，避免溢出。