第4章 TensorRT加速：模型优化与部署实践

模型训练好后，离真正的产品落地还有一条深沟。说实话，很多模型在GPU上跑得慢，不是因为显卡不够强，而是没有用对工具。TensorRT就是英伟达给深度学习推理准备的“手术刀”。

我个人习惯把TensorRT理解成一个优化器加一个运行时。它不是简单地把模型格式换一下，而是对你的网络做了一次深度整容。

4.1 为什么需要TensorRT？

你想想看，一个ResNet-50在PyTorch里跑一次推理大概需要30毫秒，用TensorRT优化后能压到5毫秒以内。这在线上服务里就是天壤之别。

它的核心优势在于：

  • 层融合：把连续的卷积、BN、激活合并成一个核（kernel）。减少了显存读写和Kernel Launch次数。
  • 精度校准：支持FP16和INT8量化。我做过一个项目，模型从FP32转到INT8，速度提升了将近3倍，精度只掉了0.3%。
  • 动态张量：允许运行时改变输入维度，不用重新构建引擎。
  • 内存复用：TensorRT会把中间张量的生命周期算得死死的，最大程度复用显存。
避坑指南：不要以为量化是万能的。我曾经在一个检测模型上直接转INT8，结果mAP从0.75掉到了0.62。后来加了校准数据集才拉回0.71。量化需要校准数据，这个步骤别省。

4.2 部署流程概览

一张图说清楚整个流程。下面这个SVG展示了从训练模型到生产服务的完整链路。

TensorRT 优化与部署全流程 训练模型 导出 ONNX 检查算子支持 构建 TensorRT 引擎 精度选择/动态形状 序列化引擎 (.engine) 加载引擎 执行推理 性能调优与验证 生产部署 关键一步 ONNX 算子必须全部被 TensorRT 支持 循环迭代直到满足 精度/延迟要求

核心路径很清晰：训练 → 导出ONNX → 构建引擎 → 部署。其中“检查算子支持”这一步，我建议你把它当铁律。一旦遇到不支持的算子，引擎构建会失败或回退到缓慢的插件模式。

4.3 第一步：导出ONNX模型

ONNX是TensorRT最稳定的输入格式。PyTorch模型导出很简单：

import torch
import torch.onnx

model = torch.load("resnet50.pth")
model.eval()

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

torch.onnx.export(
    model,
    dummy_input,
    "resnet50.onnx",
    export_params=True,
    opset_version=17,
    do_constant_folding=True,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    }
)
print("ONNX model exported successfully")

为什么设置 dynamic_axes？因为线上服务的batch size可能是动态的。这里不设的话，后面TensorRT就只能用固定batch了。

4.4 第二步：构建TensorRT引擎

我个人推荐用Python API来构建引擎，灵活度高，调优方便。

import tensorrt as trt
import numpy as np

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

def build_engine(onnx_file_path, engine_file_path, precision="fp16"):
    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 f:
        if not parser.parse(f.read()):
            for error in range(parser.num_errors):
                print(parser.get_error(error))
            return None

    config = builder.create_builder_config()
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)

    if precision == "fp16":
        if builder.platform_has_fast_fp16:
            config.set_flag(trt.BuilderFlag.FP16)
            print("FP16 mode enabled")
        else:
            print("FP16 not supported, falling back to FP32")

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

    serialized_engine = builder.build_serialized_network(network, config)
    with open(engine_file_path, "wb") as f:
        f.write(serialized_engine)

    print(f"Engine saved to {engine_file_path}")

# 调用
build_engine("resnet50.onnx", "resnet50_fp16.engine", "fp16")
配置项 说明 建议值
FP16模式 启用半精度推理，速度提升约2倍 优先使用，但需要验证精度
INT8量化 整数量化，速度提升约3-4倍 必须提供校准数据集
动态形状 支持运行时改变batch/size 设置最小、最优、最大三维
WORKSPACE 构建时临时显存上限 1GB起步，不够再加

关键经验：动态形状的“最优”维度设成生产环境最常用的尺寸。比如你线上大多数请求是batch=4，那就把(4, 3, 224, 224)设成最优。这样TensorRT的kernel选择会偏向这个维度。

4.5 第三步：运行时加载引擎

引擎文件构建一次后，后面直接加载执行。这个阶段不再需要ONNX或网络定义。

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

class TRTInference:
    def __init__(self, engine_path):
        with open(engine_path, "rb") as f:
            runtime = trt.Runtime(TRT_LOGGER)
            self.engine = runtime.deserialize_cuda_engine(f.read())
        self.context = self.engine.create_execution_context()

        # 分配显存
        self.inputs = []
        self.outputs = []
        self.bindings = []
        for i in range(self.engine.num_io_tensors):
            name = self.engine.get_tensor_name(i)
            dtype = trt.nptype(self.engine.get_tensor_dtype(name))
            shape = self.engine.get_tensor_shape(name)
            size = trt.volume(shape) * dtype().itemsize

            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
                host_buf = np.empty(shape, dtype=dtype)
                device_buf = cuda.mem_alloc(size)
                self.inputs.append({"name": name, "host": host_buf, "device": device_buf})
            else:
                host_buf = np.empty(shape, dtype=dtype)
                device_buf = cuda.mem_alloc(size)
                self.outputs.append({"name": name, "host": host_buf, "device": device_buf})

            self.bindings.append(int(device_buf))

        self.stream = cuda.Stream()

    def infer(self, input_np):
        # 设置输入形状（动态形状时需要）
        self.context.set_input_shape(self.inputs[0]["name"], input_np.shape)
        np.copyto(self.inputs[0]["host"], input_np.ravel())

        # 拷贝输入到设备
        cuda.memcpy_htod_async(self.inputs[0]["device"], self.inputs[0]["host"], self.stream)

        # 执行推理
        self.context.execute_async_v2(self.bindings, self.stream.handle)

        # 拷贝输出回主机
        cuda.memcpy_dtoh_async(self.outputs[0]["host"], self.outputs[0]["device"], self.stream)

        self.stream.synchronize()
        return self.outputs[0]["host"].reshape(self.outputs[0]["host"].shape)

# 使用
trt_infer = TRTInference("resnet50_fp16.engine")
result = trt_infer.infer(np.random.randn(4, 3, 224, 224).astype(np.float32))
print(result.shape)
注意：上面的代码用了PyCUDA做显存管理。如果你在Docker环境里跑，记得挂载--gpus all并且安装对应CUDA版本的PyCUDA。我踩过的最深的一个坑是：容器里装了两个CUDA版本，结果pycuda一直报错，查了一天才发现是环境变量乱掉了。

4.6 常见坑点与调优建议

嗯，这里必须多说几句。下面这些经验都是我用真金白银的线上事故换来的。

  1. 算子不匹配：ONNX导出时用了Torch的aten::upsample_bilinear2d，TensorRT不认识。换成torch.nn.functional.interpolate的onnx兼容版本就好了。
  2. 动态形状的min和max别太宽：如果设成(1, 3, 224, 224)到(256, 3, 224, 224)，引擎构建时会花更长时间做多profile优化，而且推理时不一定能选到最优方案。建议根据业务压测结果设定。
  3. 序列化与反序列化的CUDA版本必须一致：在TensorRT 8.x下构建的引擎，不能在TensorRT 7.x上加载。我见过直接崩溃的。
  4. 如果遇到精度敏感的场景：比如目标检测的NMS后处理，尽量保持FP32推理。我曾经试过FP16下NMS的阈值漂移导致大量漏检。
  5. 性能benchmark不可信单次结果：GPU有温度墙和频率降级，第一次跑快不代表每次都快。建议warm-up后取100次的平均延迟。
我的习惯：每次部署新模型，我会写一个自动化的精度+性能回归脚本。模型更新后，自动用旧引擎和新引擎跑同一个测试集，输出精度抖动表。这样能让线上出问题的概率降到最低。
对比项 FP32 FP16 INT8
相对速度 1x（基准） ~2x ~3-4x
精度影响 无损失 几乎无感（实验证明多数场景<0.5%） 需校准，可能掉1-3%
适用硬件 所有GPU Turing+架构（如T4、A10、A100） 需硬件INT8支持
建议场景 医疗影像、金融风控等精度第一 大部分线上服务 边缘端、高吞吐、低延迟场景

说实话，我日常工作中绝大部分模型都用FP16。INT8虽然快，但校准集的质量直接决定落地效果。如果你负责的是toC产品，用户反馈渠道又有限，建议先用FP16稳稳扛住，等模型成熟了再考虑压缩到INT8。

← 上一章 📖 返回目录 下一章 →