基于Anaconda虚拟环境的TensorRT加速推理实战：从环境配置到性能优化

1次阅读

共计 2165 个字符，预计需要花费 6 分钟才能阅读完成。

在深度学习模型的推理阶段，原生框架如 PyTorch 或 TensorFlow 往往存在以下性能瓶颈：

计算冗余：框架为训练设计的动态图机制会在推理时产生额外开销
内存占用高：默认的 FP32 精度计算对显存需求大
缺乏硬件优化：通用计算图未针对特定 GPU 架构优化

TensorRT 通过以下方式实现加速：

图优化：合并冗余计算层，消除无用分支
精度校准：支持 FP16/INT8 量化，减少计算量和内存占用
内核自动调优：为每个算子选择最优的 GPU 实现

conda create -n tensorrt_env python=3.8
conda activate tensorrt_env

关键版本匹配原则：

TensorRT 版本 → CUDA 版本 → cuDNN 版本 → PyTorch 版本
推荐组合（2023 年稳定版）：
TensorRT 8.5.x + CUDA 11.3 + cuDNN 8.6

conda install cudatoolkit=11.3 -c nvidia
conda install cudnn=8.6 -c nvidia

推荐使用 tar 包安装以便多版本共存：

tar -xzf TensorRT-8.5.3.1.Linux.x86_64-gnu.cuda-11.8.cudnn8.6.tar.gz
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$(pwd)/TensorRT-8.5.3.1/lib
pip install $(pwd)/TensorRT-8.5.3.1/python/tensorrt-8.5.3.1-cp38-none-linux_x86_64.whl

导出 ONNX 模型

torch.onnx.export(
    model,                     
    dummy_input,               
    "model.onnx",              
    opset_version=13,          # 推荐 >=13
    input_names=["input"],     
    output_names=["output"],
    dynamic_axes={"input": {0: "batch"},  # 支持动态 batch
        "output": {0: "batch"}
    }
)

转换 TensorRT 引擎

import tensorrt as trt

logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger)

# 解析 ONNX
with open("model.onnx", "rb") as f:
    parser.parse(f.read())

# 构建配置
config = builder.create_builder_config()
config.set_flag(trt.BuilderFlag.FP16)  # 启用 FP16
config.max_workspace_size = 1 << 30   # 1GB 临时内存

# 序列化引擎
serialized_engine = builder.build_serialized_network(network, config)
with open("model.engine", "wb") as f:
    f.write(serialized_engine)

FP16 量化：损失 1% 精度，速度提升 2 - 3 倍
INT8 量化：需要校准集，最大可提速 5 倍
层融合：自动合并 Conv+ReLU 等连续操作

测试环境：T4 GPU, batch_size=32

框架	延迟(ms)	吞吐量(qps)	显存占用
PyTorch FP32	45.2	708	4.3GB
TensorRT FP16	12.7	2519	1.8GB

版本冲突
现象：undefined symbol: cublasLtGetStatusString
解决：确保 conda 安装的 cudatoolkit 与系统 CUDA 版本一致
不支持的算子
自定义算子需注册插件
替代方案：用 TRT 支持的基本算子组合实现
动态 shape 处理
需在构建时指定优化 profile：

profile = builder.create_optimization_profile()
profile.set_shape("input", (1,3,224,224), (8,3,224,224), (32,3,224,224))
config.add_optimization_profile(profile)