使用TensorRT加速Pytorch模型推理

发表于 更新于

概述

TensorRT 是 NVIDIA 官方推出的深度学习推理优化工具，专为 NVIDIA GPU 设计，可显著提升模型推理速度并减少内存占用。支持 TensorFlow、PyTorch 等主流框架，通过 ONNX 实现跨框架模型转换。

一、安装TensorRT

1.1 官方包安装

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# 下载匹配CUDA版本的deb包（需NVIDIA账号）
sudo dpkg -i nv-tensorrt-repo-${os}-${tag}_1-1_amd64.deb
sudo apt-key add /var/nv-tensorrt-repo-${os}-${tag}/7fa2af80.pub
sudo apt-get update
sudo apt-get install tensorrt

# Python依赖
python3 -m pip install numpy
sudo apt-get install python3-libnvinfer-dev

# ONNX支持（可选）
python3 -m pip install numpy onnx
sudo apt-get install onnx-graphsurgeon

# PyCUDA（必需）
python3 -m pip install 'pycuda<2021.1'

1.2 验证安装

1
2
3
4
5
6
7
8
9
# 检查系统包
dpkg -l | grep TensorRT

# 验证Python环境
python3 -c "
import tensorrt
print(tensorrt.__version__)
assert tensorrt.Builder(tensorrt.Logger())
"

注意：Anaconda虚拟环境中需额外执行：

1
python3 -m pip install --upgrade nvidia-tensorrt

二、PyTorch → TensorRT 转换流程

2.1 PyTorch → ONNX

1
2
3
4
5
6
7
8
9
10
11
12
13
14
import torch
import torchvision.models as models

# 加载预训练模型
model = models.resnext50_32x4d(pretrained=True)
dummy_input = torch.randn(32, 3, 224, 224)  # BATCH_SIZE=32

# 导出ONNX
torch.onnx.export(
    model, 
    dummy_input, 
    "resnet50_onnx_model.onnx",
    verbose=False
)

替代方案：直接从 ONNX Model Zoo 下载预训练模型

2.2 ONNX → TensorRT

使用 trtexec 命令行工具（位于 /usr/src/tensorrt/bin/）：

1
2
3
4
5
6
7
8
9
10
11
# 基础转换
./trtexec --onnx=resnet50_onnx_model.onnx \
          --saveEngine=resnet_engine.trt

# 推荐参数（固定batch + FP16）
./trtexec --onnx=resnet50_onnx_model.onnx \
          --saveEngine=resnet_engine.trt \
          --explicitBatch \
          --inputIOFormats=fp16:chw \
          --outputIOFormats=fp16:chw \
          --fp16

关键参数说明：

参数 作用
--explicitBatch 固定输入batch size（与ONNX导出时一致）
--fp16 启用FP16精度（加速+省显存）
--inputIOFormats 指定输入数据格式

精度选项：TF32/FP32（默认）/FP16/INT8（需校准）

三、TensorRT 推理实现

3.1 Runtime选择

Runtime 适用场景 性能
C++ API 通用高性能需求 ★★★★
TF-TRT TensorFlow专用 ★★
Triton 多框架服务化 ★★★

推荐：PyTorch模型选用 C++ API（通过Python绑定）

3.2 Python推理代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import time

# 配置参数
BATCH_SIZE = 32
USE_FP16 = True
target_dtype = np.float16 if USE_FP16 else np.float32

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

# 内存分配
input_batch = np.random.randn(BATCH_SIZE, 224, 224, 3).astype(target_dtype)
output = np.empty([BATCH_SIZE, 1000], dtype=target_dtype)
d_input = cuda.mem_alloc(input_batch.nbytes)
d_output = cuda.mem_alloc(output.nbytes)
bindings = [int(d_input), int(d_output)]
stream = cuda.Stream()

# 预处理函数
def preprocess_input(input_data):
    # (224,224,3) → (3,224,224)
    result = torch.from_numpy(input_data).transpose(0,2).transpose(1,2)
    return np.array(result, dtype=target_dtype)

# 推理函数
def predict(batch):
    cuda.memcpy_htod_async(d_input, batch, stream)
    context.execute_async_v2(bindings, stream.handle, None)  # 异步推理
    cuda.memcpy_dtoh_async(output, d_output, stream)
    stream.synchronize()
    return output

# 执行推理
preprocessed_inputs = np.array([
    preprocess_input(inp) for inp in input_batch
])
print("Warming up...")
pred = predict(preprocessed_inputs)
print("Done warming up!")

t0 = time.time()
pred = predict(preprocessed_inputs)
print(f"Prediction cost {time.time()-t0:.4f}s")

关键点：

  • 输入/输出内存需在GPU上分配
  • 数据格式需与导出时一致（CHW vs HWC）
  • 异步推理（execute_async_v2）性能优于同步

四、性能验证

测试模型：SlowFast 行为识别
输入尺寸：(1,1,3,32,256,256)
结果：TensorRT 相比原生 PyTorch 推理速度显著提升（具体数据见原文图表）

参考资源

官方文档

社区资源