第3章模型编译与优化

模型编译与优化：从PyTorch模型到TensorRT引擎的转换流程

好，咱们今天聊聊这个环节——模型编译与优化。说白了，就是把你在PyTorch里训练好的模型，变成能在TensorRT上跑得飞快的引擎文件。这个过程，我习惯叫它“炼金术”。

你想想看，PyTorch模型就像一块璞玉，里面有各种冗余、动态图、调试信息。而TensorRT引擎，就是一把削铁如泥的宝剑。怎么把璞玉变成宝剑？这就是我们要讲的核心。

转换流程：三步走

我个人习惯把整个流程拆成三步：

导出ONNX：把PyTorch模型转成中间表示
ONNX解析与优化：TensorRT读取ONNX，做图优化
引擎构建与序列化：生成可部署的.engine文件

嗯，这里要注意，每一步都可能踩坑。我曾经在一个项目里，就因为ONNX导出时一个算子不兼容，排查了整整两天。

第一步：导出ONNX

PyTorch模型导出ONNX，核心就一个函数：torch.onnx.export()。但这里面的门道不少。

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 64, 3)
        self.relu = nn.ReLU()
    
    def forward(self, x):
        return self.relu(self.conv(x))

model = MyModel()
dummy_input = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
    opset_version=17
)

这里我特别想强调两点：

dynamic_axes：如果你需要动态batch，一定要配这个。我见过太多人忘了配，结果部署时只能固定batch size。
opset_version：建议用最新的稳定版本。太老的版本可能不支持某些算子。

避坑指南：我曾经在导出ONNX时，模型里用了torch.where()，结果ONNX不支持。最后不得不改模型结构，用torch.where(x > 0, x, 0)替代。所以，导出前最好先检查一下算子兼容性。

第二步：图优化与算子融合

这一步是TensorRT的核心竞争力。说白了，就是把计算图中的冗余操作合并、消除，让计算更高效。

我画了一张图，帮你理解这个过程：

你看，左边三个独立的Conv、BN、ReLU算子，经过融合后变成了一个算子。这有什么好处？

减少内存访问：原来每个算子都要读写一次显存，融合后只需要一次
减少kernel启动开销：GPU启动一个kernel是有开销的，算子越少，开销越小
更好的数据局部性：数据在寄存器里就能完成所有计算，不用来回搬

核心要点：TensorRT的图优化不仅仅是算子融合，还包括常量折叠、死代码消除、精度校准等。但算子融合是最直观、收益最大的优化手段。

第三步：引擎构建

这一步，TensorRT会根据你的目标GPU架构，生成最优的kernel代码。我习惯用Builder API来构建：

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_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB
config.set_flag(trt.BuilderFlag.FP16)  # 启用FP16

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

这里有几个关键点：

参数	说明	我的建议
WORKSPACE	构建时的工作空间大小	设大一点，1GB起步，否则可能构建失败
FP16	启用半精度推理	如果模型精度允许，一定要开，速度翻倍
INT8	启用INT8量化	需要校准数据集，收益更大，但调优复杂

小技巧：构建引擎时，可以指定builder.max_batch_size。如果你需要动态batch，记得在ONNX导出时就配好dynamic_axes，然后在构建时用profile设置动态范围。

避坑指南：我踩过的那些坑

做这个流程，我踩过的坑真不少。挑几个典型的说说：

ONNX算子不兼容：有一次模型里用了torch.nn.functional.interpolate，结果ONNX导出后，TensorRT解析报错。最后发现是mode参数不兼容，改成了'nearest'才解决。
动态shape问题：我刚开始做动态batch时，忘了配dynamic_axes，结果引擎只能跑固定batch size。后来重新导出，配了dynamic_axes，才搞定。
精度下降：启用FP16后，模型精度掉了0.5%。排查后发现是某个层对精度特别敏感。最后用set_precision强制该层用FP32，问题解决。

性能对比：优化前后

我拿一个实际项目的数据给你看看：

指标	PyTorch (FP32)	TensorRT (FP16)	提升
延迟 (ms)	45.2	12.8	3.5x
吞吐量 (img/s)	22	78	3.5x
显存占用 (MB)	1024	512	50%

你看，优化后的效果非常明显。延迟降了3.5倍，吞吐量翻了3.5倍，显存还省了一半。这就是TensorRT的价值所在。

嗯，整个流程就是这样。从PyTorch到ONNX，再到TensorRT引擎，每一步都有讲究。但只要你掌握了这些技巧，就能让模型跑得飞快。

📚 课程目录

模型编译与优化：从PyTorch模型到TensorRT引擎的转换流程

转换流程：三步走

第一步：导出ONNX

第二步：图优化与算子融合

第三步：引擎构建

避坑指南：我踩过的那些坑

性能对比：优化前后