当前位置： 首页 > news >正文

TensorRT详解

news 来源：原创 2025/4/29 8:58:31

文章目录

  • 一、TensorRT是什么？
  • 二、TensorRT的作用
  • 三、TensorRT的基本使用流程
  • 四、实战案例（PyTorch ➔ ONNX ➔ TensorRT 推理）
      • 1. 安装环境
      • 2. PyTorch模型导出为ONNX
      • 3. 使用TensorRT推理（Python版）
      • 案例代码
        • 案例 1：人脸检测（C++）
        • 案例 2：UNet 医学影像分割（Python）
      • 优化策略与注意事项
  • 五、TensorRT常用高级功能
  • 六、总结
    • 应用场景
  • 进阶应用案例
    • 动态尺寸推理（Dynamic Shape）
      • 1. 动态尺寸构建流程
        • 示例代码（稍复杂）
    • INT8量化推理
      • 2. INT8量化推理的构建流程
        • 示例代码（简单版）
    • 小结一下：
    • 总结TensorRT部署技巧大全

一、TensorRT是什么？

TensorRT（Tensor Runtime） 是 NVIDIA 开发的一个高性能深度学习推理（inference）优化库和运行时库，它可以帮助你将训练好的深度学习模型加速部署到 NVIDIA GPU 上，专注于在 NVIDIA GPU 上实现低延迟、高吞吐量的模型部署。

  • 核心定位：神经网络推理优化和加速。
  • 支持模型：支持 TensorFlow、PyTorch、ONNX 等训练框架导出的模型。
  • 主要功能：
    • 网络结构优化（模型优化）：（比如图融合、层合并），通过算子融合（如卷积、BN、激活层合并）、消除冗余计算（如删除无用层）重构计算图，减少内存占用和计算量。
    • 精度削减（精度校准）（FP32 ➔ FP16 或 INT8）：支持 FP32/FP16/INT8 等量化技术，在精度损失可控的前提下提升速度（如 INT8 相比 FP32 速度提升 4 倍）。
    • 内存/带宽优化
    • 并行计算和张量融合
    • 动态张量调整（dynamic shapes）
    • 硬件适配：自动选择适合 GPU 架构的最优内核，动态管理显存，并支持多流并行处理任务
    • 跨框架兼容：支持 TensorFlow、PyTorch（通过 ONNX）、Caffe 等框架模型的转换与优化

TensorRT在服务器、边缘设备（如Jetson系列）、甚至自动驾驶芯片（如Drive PX）上被广泛应用。

二、TensorRT的作用

方面说明
推理加速大幅度降低推理延迟，提升吞吐量。通常能加速 2-10 倍。
资源节省通过量化（FP16/INT8）减少模型大小，降低显存占用。
多平台支持支持服务器、嵌入式设备（如Jetson Nano、Xavier）等。
灵活性强支持动态输入尺寸（Dynamic Shape）、序列化部署（Engine文件）。
提升吞吐量通过动态批处理和多流执行，推荐系统吞吐量可提升 6 倍。
节能高效适用于自动驾驶、工业检测等实时性要求高的场景。

三、TensorRT的基本使用流程

大致步骤如下：

  1. 模型导出：将训练好的模型转换成 ONNX 格式。
# PyTorch 示例  
import torch  
model = torch.hub.load('mateuszbuda/brain-segmentation-pytorch', 'unet', pretrained=True)  
dummy_input = torch.randn(1, 3, 256, 256)  
torch.onnx.export(model, dummy_input, "unet.onnx", input_names=["input"], output_names=["output"])
  1. 模型解析：用 TensorRT 的 ONNX Parser 读取模型。
  2. 构建TensorRT优化引擎（Engine）：选择精度（FP32、FP16、INT8）并生成优化推理引擎。
// C++ 示例（网页2）  
#include <NvInfer.h>  
using namespace nvinfer1;  std::unique_ptr<ICudaEngine> initEngine(const std::string& onnx_path) {  IBuilder* builder = createInferBuilder(gLogger);  INetworkDefinition* network = builder->createNetwork();  OnnxParser* parser = createOnnxParser(network);  parser->parseFromFile(onnx_path.c_str(), static_cast<int>(ILogger::Severity::kWARNING));  builder->setMaxBatchSize(1);  builder->setMaxWorkspaceSize(1 << 20);  return std::unique_ptr<ICudaEngine>(builder->buildCudaEngine(network));  
}
  1. 推理执行（Inference）：用引擎推理输入数据。
import tensorrt as trt  
import pycuda.driver as cuda  # 加载引擎  
with open("model.engine", "rb") as f:  engine = runtime.deserialize_cuda_engine(f.read())  
context = engine.create_execution_context()  # 分配内存并推理  
h_input = cuda.pagelocked_empty(trt.volume(context.get_binding_shape(0)), dtype=np.float32)  
d_input = cuda.mem_alloc(h_input.nbytes)  
stream = cuda.Stream()  
cuda.memcpy_htod_async(d_input, h_input, stream)  
context.execute_async_v2(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)

  • （可选）序列化保存引擎：下次直接加载，不需要重新构建。

  • 关键配置

  • 动态输入：通过 setMaxBatchSize 和 setMaxWorkspaceSize 调整批次和显存。

  • 混合精度：设置 precision_mode 为 FP16 或 INT8，需校准数据集以减少精度损失。

四、实战案例（PyTorch ➔ ONNX ➔ TensorRT 推理）

1. 安装环境

# 安装TensorRT（需要有NVIDIA GPU）
# 在官网根据你的CUDA版本下载对应的TensorRT包
# 官方地址：https://developer.nvidia.com/tensorrt# ONNX工具
pip install onnx onnxruntime onnx-simplifier
pip install torch2trt  # 可选，用于快速PyTorch转TensorRT

2. PyTorch模型导出为ONNX

假设有一个简单的PyTorch模型：

import torch
import torch.nn as nnclass SimpleModel(nn.Module):def __init__(self):super(SimpleModel, self).__init__()self.fc = nn.Linear(10, 5)def forward(self, x):return self.fc(x)model = SimpleModel()
dummy_input = torch.randn(1, 10)# 导出为ONNX格式
torch.onnx.export(model, dummy_input, "simple_model.onnx", input_names=["input"], output_names=["output"])

3. 使用TensorRT推理（Python版）

import tensorrt as trt
import numpy as np
import pycuda.driver as cuda
import pycuda.autoinitTRT_LOGGER = trt.Logger(trt.Logger.WARNING)# 1. 读取ONNX模型并构建引擎
def build_engine(onnx_file_path):builder = trt.Builder(TRT_LOGGER)network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))parser = trt.OnnxParser(network, TRT_LOGGER)with open(onnx_file_path, 'rb') as model:if not parser.parse(model.read()):for error in range(parser.num_errors):print(parser.get_error(error))return Nonebuilder.max_workspace_size = 1 << 20  # 1MBbuilder.fp16_mode = True  # 开启FP16加速（需要硬件支持）engine = builder.build_cuda_engine(network)return engine# 2. 加载引擎并推理
def inference(engine, input_data):context = engine.create_execution_context()# 准备输入输出input_shape = (1, 10)output_shape = (1, 5)# 分配GPU内存d_input = cuda.mem_alloc(input_data.nbytes)d_output = cuda.mem_alloc(np.empty(output_shape, dtype=np.float32).nbytes)bindings = [int(d_input), int(d_output)]stream = cuda.Stream()# 拷贝输入到GPUcuda.memcpy_htod_async(d_input, input_data, stream)# 推理context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)# 从GPU拷回结果output_data = np.empty(output_shape, dtype=np.float32)cuda.memcpy_dtoh_async(output_data, d_output, stream)stream.synchronize()return output_data# 使用示例
engine = build_engine("simple_model.onnx")
input_data = np.random.randn(1, 10).astype(np.float32)
output = inference(engine, input_data)
print(output)

案例代码

案例 1：人脸检测（C++）
std::vector<cv::Rect> detectFaces(cv::Mat& image, std::unique_ptr<ICudaEngine>& engine) {  IExecutionContext* context = engine->createExecutionContext();  float* inputData = preprocess(image); // 预处理图像  void* bindings[] = { inputData };  context->executeV2(bindings);  float* detectionOutput = new float[1 * 100 * 7];  context->getOutputHandle(0)->copyToHost(detectionOutput);  // 解析检测结果...  
}
案例 2：UNet 医学影像分割（Python）
# 加载 ONNX 模型并生成引擎  
import tensorrt as trt  
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)  
with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.OnnxParser(network, TRT_LOGGER) as parser:  with open("unet.onnx", "rb") as model:  parser.parse(model.read())  config = builder.create_builder_config()  config.max_workspace_size = 1 << 30  engine = builder.build_engine(network, config)

优化策略与注意事项

  1. 性能调优：
    • 调整 max_workspace_size 以允许更多优化算法。
    • 使用动态批处理提升吞吐量。
  2. 硬件兼容性：
    • 引擎与编译时的 GPU 架构绑定，跨代硬件需重新优化。
  3. 常见问题：
    • INT8 精度损失：采用量化感知训练（QAT）而非后训练量化（PTQ）。
    • 内存溢出：增大工作空间或启用显存池复用。

五、TensorRT常用高级功能

功能说明
FP16精度推理只需要设置builder.fp16_mode = True
INT8量化推理需要配合校准数据集使用，进行量化校准（更高效，但需要仔细处理精度问题）
动态尺寸支持支持变形输入，比如视频帧不同分辨率
网络层融合自动优化多个算子，减少推理延迟
序列化引擎推理时直接加载.plan文件，跳过ONNX解析

六、总结

  • TensorRT 主要用于推理加速，不是用来训练模型的。
  • 最重要的是熟悉 模型导出 → 优化构建 → 推理执行 这一套流程。
  • 在服务器部署（如TensorRT-Server / Triton Inference Server）和边缘设备部署（如Jetson）时，TensorRT都是必不可少的。
  • 注意TensorRT版本与CUDA/cuDNN兼容性，一般官网文档有详细说明。

应用场景

  • 自动驾驶：YOLOv5 目标检测优化后帧率提升至 200FPS。
  • 医疗影像：肿瘤检测模型推理时间降至 6.14ms，支持实时诊断。
  • 自然语言处理：BERT-Large 推理加速至 1.2ms，支持实时交互。

通过以上步骤和技术，TensorRT 能够显著提升模型推理效率，适用于从嵌入式设备到数据中心的多样化场景。更多实现细节可参考 NVIDIA 官方文档和示例代码库。

好的！
既然你想深入了解，我们继续讲一下 TensorRT 动态尺寸推理（Dynamic Shape） 和 INT8量化加速 —— 这两个是实际部署中很重要的高级技巧。

进阶应用案例

动态尺寸推理（Dynamic Shape）

为什么要动态尺寸？
在实际应用中，输入数据尺寸可能是变化的，比如：

  • 视频流中的图像大小不同
  • 语音识别中语音片段长度不同
  • 医学图像中不同分辨率扫描件

传统TensorRT默认是固定输入尺寸，如果要处理变尺寸输入，就需要开启动态尺寸（Optimization Profile）！

1. 动态尺寸构建流程

需要增加：

  • optimization_profile：告诉TensorRT哪些输入尺寸范围是合法的。
  • 在推理时，根据实际输入尺寸重新设置binding。
示例代码（稍复杂）
import tensorrt as trt
import numpy as np
import pycuda.driver as cuda
import pycuda.autoinitTRT_LOGGER = trt.Logger(trt.Logger.WARNING)def build_engine_with_dynamic_shape(onnx_file_path):builder = trt.Builder(TRT_LOGGER)network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))parser = trt.OnnxParser(network, TRT_LOGGER)with open(onnx_file_path, 'rb') as model:parser.parse(model.read())config = builder.create_builder_config()config.max_workspace_size = 1 << 20  # 1MB# 开启FP16（可选）config.set_flag(trt.BuilderFlag.FP16)# 设置动态shapeprofile = builder.create_optimization_profile()input_name = network.get_input(0).name# 设置最小、最优、最大尺寸（batch_size, dim1, dim2）profile.set_shape(input_name, (1, 10), (4, 10), (8, 10))config.add_optimization_profile(profile)engine = builder.build_engine(network, config)return enginedef inference_dynamic(engine, input_data):context = engine.create_execution_context()# 必须设置实际输入尺寸context.set_binding_shape(0, input_data.shape)# 分配GPU内存d_input = cuda.mem_alloc(input_data.nbytes)output_shape = (input_data.shape[0], 5)d_output = cuda.mem_alloc(np.empty(output_shape, dtype=np.float32).nbytes)bindings = [int(d_input), int(d_output)]stream = cuda.Stream()cuda.memcpy_htod_async(d_input, input_data, stream)context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)output_data = np.empty(output_shape, dtype=np.float32)cuda.memcpy_dtoh_async(output_data, d_output, stream)stream.synchronize()return output_data# 使用示例
engine = build_engine_with_dynamic_shape("simple_model.onnx")
input_data = np.random.randn(3, 10).astype(np.float32)  # 动态batch=3
output = inference_dynamic(engine, input_data)
print(output)

INT8量化推理

为什么用INT8？

  • 比FP32快2-4倍
  • 模型占用显存少，适合边缘设备
  • 小心精度丢失，需要"校准"！

2. INT8量化推理的构建流程

TensorRT需要一个校准器，用一小批真实输入数据来估计量化的 scale/zero-point。

可以使用：

  • 内置的EntropyCalibrator
  • 自定义一个Calibrator类（更灵活）
示例代码（简单版）
import tensorrt as trt
import numpy as np
import pycuda.driver as cuda
import pycuda.autoinit
import randomTRT_LOGGER = trt.Logger(trt.Logger.WARNING)# 1. 自定义校准器
class MyCalibrator(trt.IInt8EntropyCalibrator2):def __init__(self, calibration_data):super(MyCalibrator, self).__init__()self.data = calibration_dataself.index = 0self.device_input = cuda.mem_alloc(self.data[0].nbytes)def get_batch_size(self):return 1def get_batch(self, names):if self.index >= len(self.data):return Nonecuda.memcpy_htod(self.device_input, self.data[self.index])self.index += 1return [int(self.device_input)]def read_calibration_cache(self):return Nonedef write_calibration_cache(self, cache):pass# 2. 构建引擎
def build_engine_int8(onnx_file_path, calibration_data):builder = trt.Builder(TRT_LOGGER)network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))parser = trt.OnnxParser(network, TRT_LOGGER)with open(onnx_file_path, 'rb') as model:parser.parse(model.read())config = builder.create_builder_config()config.max_workspace_size = 1 << 20config.set_flag(trt.BuilderFlag.INT8)calibrator = MyCalibrator(calibration_data)config.int8_calibrator = calibratorengine = builder.build_engine(network, config)return engine# 使用示例
calibration_data = [np.random.randn(1, 10).astype(np.float32) for _ in range(10)]  # 模拟10组校准数据
engine = build_engine_int8("simple_model.onnx", calibration_data)# 后续推理方法与之前一样

小结一下：

技巧核心点注意事项
动态尺寸推理设置Optimization Profile + 动态Binding推理时一定要设置实际输入尺寸
INT8量化推理配置Int8Calibrator校准器量化后精度可能有下降，需要评估

总结TensorRT部署技巧大全

  • 动态输入：profile.set_shape
  • FP16推理：config.set_flag(trt.BuilderFlag.FP16)
  • INT8推理：config.set_flag(trt.BuilderFlag.INT8)
  • 多个profile：支持不同输入分辨率（比如640×480，1280×720）
  • 引擎序列化：加速加载，避免每次重建
  • 与TensorFlow、PyTorch无缝对接（通过ONNX中间件）

相关文章：

  • 练习普通话，说话更有节奏
  • Matplotlib可视化基础
  • Transformer 原理逐行解析：从 Self-Attention 到 Positional Encoding
  • DeepSeek-R1技术报告（中文版）
  • 质量的“试金石”：精通Spring Boot单元测试与集成测试
  • 简单理解https与http
  • GESP2024年9月认证C++八级( 第二部分判断题（6-10））
  • WSL释放空间
  • JavaScript性能优化实战（6）：网络请求与资源加载优化
  • 【刷题Day29】Python/JAVA - 03（浅）
  • CAD编程的知识
  • 什么是 DDoS 攻击？高防 IP 如何有效防护？2025全面解析与方案推荐
  • terraform使用workspace管理多工作环境
  • 一文掌握Matplotlib绘图
  • 【Kubernetes】部署 Kubernetes 仪表板（Dashboard）
  • 《Linux篇》基础开发工具——vim详细介绍
  • Nacos-3.0.0适配PostgreSQL数据库
  • CUDA 编程相关的开源库
  • 单片机-89C51部分：6、数码管
  • 基于卷积神经网络的蔬菜水果识别系统，resnet50，mobilenet模型【pytorch框架+python源码】
  • 一位排球青训教练的20年时光：努力提高女排球员成才率
  • 论法的精神︱张玉敏：知识产权保护要为社会经济文化发展服务
  • 涨价应对关税变化是短期之策，跨境电商塑造新品牌开辟“新蓝海”
  • 钟声：美以芬太尼为借口滥施关税，纯属“内病外治”
  • 第二十届华表奖提名名单公布，张译、王一博、马丽、郭帆等入围
  • 苏迪曼杯即将在厦门打响，国羽向创纪录的14冠进军