Day 43：TensorRT实战 — 模型转换与优化

今日目标 ：掌握PyTorch→ONNX→TensorRT的完整部署流水线，学会FP16/INT8引擎构建，理解TensorRT的优化原理

预计阅读 ：12分钟 | 动手操作 ：40分钟

前置条件：需要NVIDIA GPU + TensorRT环境（Jetson或PC均可）

一、TensorRT是什么？

复制代码

TensorRT：NVIDIA的深度学习推理优化器

核心功能：
  1. 模型转换：PyTorch/TF → TensorRT Engine
  2. 精度优化：FP32 → FP16 → INT8
  3. 图优化：层融合、消除冗余操作
  4. 内存优化：显存复用、减少分配

为什么TensorRT这么快？
  1. 层融合 (Layer Fusion)
     Conv + BN + ReLU → 融合成一个CBR核
     减少显存读写，提升吞吐量
     
  2. 精度校准 (Precision Calibration)
     FP32 → FP16/INT8，保持精度
     
  3. 核自动调优 (Kernel Auto-Tuning)
     针对当前GPU选择最优的CUDA kernel
     
  4. 动态显存管理
     复用显存，减少分配开销
     
  5. 多流并行
     同时执行多个操作

TensorRT的典型加速比：
  FP32: 1-2x (图优化)
  FP16: 2-4x (Tensor Core)
  INT8: 3-6x (DP4A指令)

python 复制代码

import torch
import torch.nn as nn
import torch.onnx
import numpy as np
import os

torch.manual_seed(42)

二、核心知识点

2.1 TensorRT优化：层融合

python 复制代码

"""
层融合是TensorRT最重要的优化之一

融合前：
  Conv → Bias → ReLU
  三个独立操作，每个都需要读写显存
  
  显存读写：
    Conv: 读输入 → 写输出
    Bias: 读Conv输出 → 写Bias输出
    ReLU: 读Bias输出 → 写最终输出
    = 3次读 + 3次写

融合后 (CBR):
  Conv + Bias + ReLU → 一个融合核
  显存读写：
    读输入 → 写输出
    = 1次读 + 1次写

显存带宽节省：67%！
对于显存带宽敏感的端侧设备（如Jetson），这是巨大的提升

TensorRT支持的融合模式：
  Conv + Bias + ReLU → CBR
  Conv + Bias + Add + ReLU → CBR + 残差
  Conv + BN + ReLU → CBR (BN参数fold到Conv中)
  Conv + SiLU → CBR
  MatMul + Add + ReLU → 全连接层融合
"""

def layer_fusion_demo():
    print("TensorRT层融合:")
    print()
    print("融合前 (3个操作):")
    print("  Conv → [读输入, 写输出]")
    print("  Bias → [读中间, 写中间]")
    print("  ReLU → [读中间, 写输出]")
    print("  总计: 3次读 + 3次写")
    print()
    print("融合后 (1个操作):")
    print("  CBR → [读输入, 写输出]")
    print("  总计: 1次读 + 1次写")
    print("  节省: 67%显存带宽")
    print()
    print("端侧AI (Jetson):")
    print("  Jetson的显存带宽是瓶颈")
    print("  层融合直接提升速度 1.5-2x")

layer_fusion_demo()

2.2 PyTorch → ONNX：模型导出

python 复制代码

"""
ONNX (Open Neural Network Exchange)
  开放神经网络交换格式
  是PyTorch → TensorRT的桥梁

为什么需要ONNX？
  PyTorch是动态图 → TensorRT需要静态图
  ONNX是静态图 → 适合TensorRT
  而且ONNX是跨框架的：TF/Caffe/PyTorch都可以导出

导出ONNX的关键参数：
  input_names: 输入节点名称
  output_names: 输出节点名称
  opset_version: ONNX算子集版本 (推荐11+)
  dynamic_axes: 动态维度 (batch可变)
  export_params: 是否导出参数 (True)
"""

class SimpleModel(nn.Module):
    """用于演示导出的简单模型"""
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(64, 128, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(128)
        self.fc = nn.Linear(128 * 16 * 16, 10)
    
    def forward(self, x):
        x = self.relu(self.bn1(self.conv1(x)))
        x = self.pool(x)
        x = self.relu(self.bn2(self.conv2(x)))
        x = self.pool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x


def export_to_onnx_demo():
    """演示PyTorch模型导出ONNX"""
    model = SimpleModel()
    model.eval()
    
    # 创建虚拟输入
    dummy_input = torch.randn(1, 3, 32, 32)
    
    # 导出ONNX
    onnx_path = "simple_model.onnx"
    
    try:
        torch.onnx.export(
            model,
            dummy_input,
            onnx_path,
            export_params=True,        # 导出参数
            opset_version=11,          # ONNX算子集版本
            do_constant_folding=True,  # 常量折叠优化
            input_names=['input'],     # 输入节点名
            output_names=['output'],   # 输出节点名
            dynamic_axes={             # 动态维度
                'input': {0: 'batch_size'},
                'output': {0: 'batch_size'},
            },
        )
        
        print("ONNX导出成功!")
        print(f"  模型: {onnx_path}")
        print(f"  大小: {os.path.getsize(onnx_path) / 1024:.1f} KB")
        print()
        print("ONNX导出关键参数:")
        print("  opset_version=11: 推荐，支持大部分算子")
        print("  do_constant_folding=True: 预先计算常量")
        print("  dynamic_axes: batch_size可变")
        
    except Exception as e:
        print(f"ONNX导出失败: {e}")
        print("(这是正常的，需要安装onnx库)")
    
    # 清理
    if os.path.exists(onnx_path):
        os.remove(onnx_path)


export_to_onnx_demo()

2.3 ONNX → TensorRT Engine

python 复制代码

"""
TensorRT Engine构建流程

有两种方式：
1. 使用trtexec命令行工具（推荐）
2. 使用Python API

方式1: trtexec (命令行)
  FP32:
    trtexec --onnx=model.onnx --saveEngine=model_fp32.engine
  
  FP16:
    trtexec --onnx=model.onnx --saveEngine=model_fp16.engine --fp16
  
  INT8:
    trtexec --onnx=model.onnx --saveEngine=model_int8.engine --int8

方式2: Python API
  import tensorrt as trt
  
  builder = trt.Builder(logger)
  network = builder.create_network()
  parser = trt.OnnxParser(network, logger)
  parser.parse_from_file('model.onnx')
  
  config = builder.create_builder_config()
  config.set_flag(trt.BuilderFlag.FP16)  # 启用FP16
  
  engine = builder.build_engine(network, config)

构建Engine的时间：
  小模型(ResNet): ~1分钟
  中等模型(YOLO): ~5分钟
  大模型(ViT): ~10-30分钟

Engine可以序列化保存，推理时直接加载
"""

def tensorrt_engine_build_demo():
    print("TensorRT Engine构建流程:")
    print()
    print("方式1: trtexec (命令行，推荐)")
    print("  FP32: trtexec --onnx=model.onnx --saveEngine=model_fp32.engine")
    print("  FP16: trtexec --onnx=model.onnx --saveEngine=model_fp16.engine --fp16")
    print("  INT8: trtexec --onnx=model.onnx --saveEngine=model_int8.engine --int8")
    print()
    print("方式2: Python API")
    print("  1. 创建builder")
    print("  2. 创建network")
    print("  3. 解析ONNX")
    print("  4. 配置精度 (FP16/INT8)")
    print("  5. build engine")
    print()
    print("构建时间:")
    print("  ResNet-50: ~1分钟")
    print("  YOLOv8: ~5分钟")
    print("  ViT: ~10-30分钟")
    print()
    print("Engine文件:")
    print("  model_fp32.engine: ~100MB")
    print("  model_fp16.engine: ~50MB")
    print("  model_int8.engine: ~25MB")

tensorrt_engine_build_demo()

2.4 TensorRT推理流程

python 复制代码

"""
TensorRT推理流程

1. 加载Engine
2. 创建ExecutionContext
3. 分配输入输出显存
4. 拷贝输入数据到GPU
5. 执行推理 (execute_v2)
6. 拷贝输出数据到CPU
7. 后处理

TensorRT推理代码模板：
"""

def tensorrt_inference_template():
    code = """
# TensorRT推理模板
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit

# 1. 加载Engine
with open('model_fp16.engine', 'rb') as f:
    engine_data = f.read()

runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
engine = runtime.deserialize_cuda_engine(engine_data)

# 2. 创建ExecutionContext
context = engine.create_execution_context()

# 3. 分配显存
input_shape = (1, 3, 224, 224)
output_shape = (1, 1000)

d_input = cuda.mem_alloc(int(np.prod(input_shape) * 4))
d_output = cuda.mem_alloc(int(np.prod(output_shape) * 4))

# 4. 拷贝输入到GPU
input_data = np.random.randn(*input_shape).astype(np.float32)
cuda.memcpy_htod(d_input, input_data)

# 5. 执行推理
bindings = [int(d_input), int(d_output)]
context.execute_v2(bindings)

# 6. 拷贝输出到CPU
output_data = np.empty(output_shape, dtype=np.float32)
cuda.memcpy_dtoh(output_data, d_output)

# 7. 后处理
pred_class = np.argmax(output_data)
print(f"预测类别: {pred_class}")
"""
    print("TensorRT推理模板:")
    print(code)

tensorrt_inference_template()

三、动手实践：对比三种精度

python 复制代码

def precision_comparison():
    """
    FP32 / FP16 / INT8 三种精度对比
    """
    comparison = {
        'FP32': {
            '模型大小': '100MB (基准)',
            '推理速度': '1x (基准)',
            '精度': '基准',
            '显存占用': '基准',
            '构建时间': '无需校准',
            'build flag': '默认',
        },
        'FP16': {
            '模型大小': '50MB (50%)',
            '推理速度': '2-4x',
            '精度': '≈ FP32 (几乎无损)',
            '显存占用': '50%',
            '构建时间': '无需校准',
            'build flag': '--fp16',
        },
        'INT8': {
            '模型大小': '25MB (25%)',
            '推理速度': '3-6x',
            '精度': '≈ FP32 (PTQ) / = FP32 (QAT)',
            '显存占用': '25%',
            '构建时间': '需要校准(500-1000张图)',
            'build flag': '--int8',
        },
    }
    
    print("三种精度对比:")
    for precision, info in comparison.items():
        print(f"\n  [{precision}]")
        for k, v in info.items():
            print(f"    {k}: {v}")
    
    print("\n\n端侧AI推荐:")
    print("  Jetson Orin/Nano: FP16 (简单，精度高)")
    print("  Jetson Xavier: INT8 (追求极致性能)")
    print("  桌面GPU: FP16 (日常开发)")

precision_comparison()

四、端侧AI部署流水线

python 复制代码

def edge_deployment_pipeline():
    """
    端侧AI完整部署流水线
    """
    pipeline = {
        'Step 1': {
            '操作': 'PyTorch模型训练',
            '工具': 'PyTorch',
            '输出': 'model.pth',
        },
        'Step 2': {
            '操作': '导出ONNX',
            '工具': 'torch.onnx.export',
            '输出': 'model.onnx',
        },
        'Step 3': {
            '操作': 'ONNX模型检查',
            '工具': 'onnx.checker + netron',
            '输出': '验证通过',
        },
        'Step 4': {
            '操作': 'TensorRT Engine构建',
            '工具': 'trtexec 或 Python API',
            '输出': 'model_fp16.engine',
        },
        'Step 5': {
            '操作': '精度验证',
            '工具': '对比PyTorch和TensorRT输出',
            '输出': '精度报告',
        },
        'Step 6': {
            '操作': '端侧推理部署',
            '工具': 'TensorRT Runtime + PyCUDA',
            '输出': '推理结果',
        },
    }
    
    print("端侧AI部署流水线 (PyTorch → TensorRT):")
    for step, info in pipeline.items():
        print(f"\n  {step}: {info['操作']}")
        print(f"    工具: {info['工具']}")
        print(f"    输出: {info['输出']}")

edge_deployment_pipeline()

五、常见坑点

坑1：ONNX导出时算子不支持

python 复制代码

# 不是所有PyTorch算子都能导出ONNX
# 常见不支持的操作：
# - torch.einsum (部分支持)
# - 自定义CUDA kernel
# - 动态控制流 (if/while)

# 排查方法：
# 1. 检查opset_version是否够新
# 2. 用onnx.checker检查模型
# 3. 用netron可视化，看是否有未知节点
# 4. 对于不支持的操作，改写为支持的等价操作

# 常见替代：
# torch.einsum → torch.bmm + permute
# 自定义操作 → 注册为ONNX custom op

坑2：TensorRT版本和ONNX opset不匹配

python 复制代码

# TensorRT版本和ONNX opset版本需要匹配
# TensorRT 8.0: 支持到opset 13
# TensorRT 8.4: 支持到opset 14
# TensorRT 8.6: 支持到opset 17

# 如果opsnet版本太高，TensorRT不支持
# 解决方案：降低opset_version
# torch.onnx.export(..., opset_version=11)  # 安全选择

坑3：INT8校准数据的选择

python 复制代码

# INT8校准需要代表性数据
# 校准数据应该：
# 1. 覆盖推理时的数据分布
# 2. 数量足够 (500-1000张)
# 3. 多样化 (各种场景、光照、角度)

# ❌ 错误：只用10张图片校准
# → 量化scale不准确 → 精度大幅下降

# ✅ 正确：用500+张代表性图片
# → 量化scale准确 → 精度下降 < 1%

坑4：动态shape的问题

python 复制代码

# TensorRT默认不支持动态shape
# 如果输入尺寸不固定，需要配置

# 解决方案1：固定尺寸（最简单）
# 所有输入都resize到固定尺寸

# 解决方案2：动态shape（复杂）
# 在build时配置optimization profile
# 指定min/opt/max三种尺寸
# builder_profile.set_shape("input", min=(1,3,224,224), 
#                                       opt=(1,3,512,512), 
#                                       max=(1,3,1024,1024))

六、今日作业

导出ONNX：把之前训练的模型导出为ONNX，用netron可视化
TensorRT探索：如果有GPU，用trtexec构建FP16/INT8 engine
对比速度：对比PyTorch和TensorRT的推理速度
打卡：评论区发你的理解，格式："Day 43/100 打卡：TensorRT已掌握！"

今日小结

复制代码

今天你学会了：
✅ TensorRT的五大优化：层融合/精度校准/核调优/显存管理/多流
✅ 层融合：Conv+BN+ReLU → 1个操作，节省67%显存带宽
✅ PyTorch → ONNX：导出参数和关键配置
✅ ONNX → TensorRT Engine：trtexec + Python API
✅ FP32/FP16/INT8三种精度对比
✅ INT8校准：需要500-1000张代表性图片
✅ 端侧AI部署流水线：6步从PyTorch到端侧推理
✅ 4个经典坑点

明日预告

Day 44：模型剪枝与稀疏化

结构化剪枝、非结构化剪枝、彩票假说、端侧剪枝实践

🔥 关注我，每天解锁一个端侧AI技能！

微信公众号：xxx | 小红书：xxx | CSDN：xxx

评论区打卡，一起坚持100天！

附：小红书图文版

封面标题建议：TensorRT实战 | 模型加速2-6倍，端侧部署必备！🚀

P1 --- 封面

标题：TensorRT实战

副标题：ONNX导出 / Engine构建 / FP16/INT8

P2 --- TensorRT五大优化

层融合：Conv+BN+ReLU→1个操作

精度校准：FP32→FP16/INT8

核自动调优：选最优CUDA kernel

动态显存管理：复用显存

多流并行：同时执行多个操作

P3 --- 部署流水线

PyTorch → ONNX → TensorRT

训练模型 → model.pth
导出ONNX → model.onnx (netron检查)
构建Engine → model_fp16.engine
端侧推理 → 2-6x加速！

P4 --- 三种精度对比

FP32: 100MB, 1x速度, 基准精度

FP16: 50MB, 2-4x, 几乎无损

INT8: 25MB, 3-6x, 需要校准

端侧推荐：FP16 (简单，精度高)

P5 --- 层融合的秘密

Conv+Bias+ReLU → CBR

3次读写 → 1次读写

节省67%显存带宽

Jetson的显存带宽是瓶颈

层融合直接提速1.5-2x

P6 --- 常见坑点

ONNX算子不支持？降opset或用等价操作

版本不匹配？TensorRT 8.0→opset 11

INT8精度差？校准数据500+张

动态shape？配置optimization profile

评论区打卡 Day 43/100

标签：#TensorRT #ONNX #模型部署 #端侧AI #Jetson

CSDN发布提示：CSDN版本建议在部署流水线部分放完整的PyTorch→ONNX→TensorRT流程图，在层融合部分放融合前后的对比图，在三种精度对比部分放对比表，在INT8校准部分放校准流程图。