模型导出与转换：PyTorch 转 ONNX，再到 TensorRT Engine

模型训练好了，只是第一步。在工业相机边上那个小盒子（通常是 Jetson 或者 x86 工控机）里跑起来，才是真本事。这一步，说白了就是让模型「换个语言」——从 PyTorch 的 Python 世界，搬到 TensorRT 的 C++ 引擎里。

我个人习惯把这条链路拆成三段：PyTorch → ONNX，ONNX → TensorRT Engine，最后是精度校准。每一段都有坑，我踩过不少，今天一并倒出来。

4.1 PyTorch 模型导出为 ONNX

ONNX 是个中间格式，像个翻译官。PyTorch 自带的 torch.onnx.export 就能干这事。但别以为一行代码就搞定，参数调不对，后面 TensorRT 解析时会哭。

举个例子，我做过一个钢材表面缺陷检测模型，输入是 [1, 3, 640, 640]。但实际产线来的图不一定是 640x640，需要动态宽高。怎么办？

import torch

model = YourModel().eval()
dummy_input = torch.randn(1, 3, 640, 640)

torch.onnx.export(
    model,
    dummy_input,
    "defect_model.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch_size", 2: "height", 3: "width"},
        "output": {0: "batch_size"}
    },
    opset_version=11,
    do_constant_folding=True
)

这里我特别强调 opset_version。我建议用 11 或更高。为什么？低版本 opset 不支持某些算子，比如 Resize 在 opset 10 里行为跟 11 不一样。我曾经因为 opset 设成 9，导出的 ONNX 在 TensorRT 里直接崩了，查了两天才发现是版本不匹配。

4.2 ONNX 转 TensorRT Engine

拿到 ONNX 文件，下一步就是用 trtexec 或者 TensorRT Python API 转成 Engine。Engine 是 TensorRT 的运行时格式，针对你的目标 GPU 做了极致优化。

命令行方式最直接，适合批量测试：

trtexec --onnx=defect_model.onnx \
        --saveEngine=defect_model.engine \
        --fp16 \
        --workspace=4096 \
        --minShapes=input:1x3x320x320 \
        --optShapes=input:1x3x640x640 \
        --maxShapes=input:1x3x1280x1280

注意 --minShapes、--optShapes、--maxShapes 这三个参数。它们定义了动态 shape 的范围。TensorRT 会在优化时选一个最优 shape（就是 optShapes）做 kernel 调优。如果你的实际输入 shape 离 optShapes 太远，性能会打折扣。

用 Python API 转的好处是可以嵌入到自动化流程里：

import tensorrt as trt

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

with open("defect_model.onnx", "rb") as f:
    parser.parse(f.read())

config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4 << 30)  # 4GB
config.set_flag(trt.BuilderFlag.FP16)

profile = builder.create_optimization_profile()
profile.set_shape("input", (1, 3, 320, 320), (1, 3, 640, 640), (1, 3, 1280, 1280))
config.add_optimization_profile(profile)

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

嗯，这里要注意：build_serialized_network 返回的是序列化后的 engine，可以直接存文件。TensorRT 8.5 之后推荐用这个接口，老版本用 build_engine 也行，但序列化方式不同。

4.3 精度校准：FP16 与 INT8

模型转成 FP16 一般损失不大，但 INT8 就敏感多了。工业缺陷检测里，一个微小的裂纹可能只占几个像素，INT8 量化不当，这些细节就丢了。

FP16 转换很简单，上面代码里已经展示了——加个 set_flag(trt.BuilderFlag.FP16) 就行。但 INT8 需要校准数据集。

TensorRT 提供了 Int8Calibrator 接口，你需要继承它并实现 get_batch 方法：

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

class DefectCalibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, calibration_files, batch_size, input_shape):
        trt.IInt8EntropyCalibrator2.__init__(self)
        self.batch_size = batch_size
        self.input_shape = input_shape
        self.calibration_files = calibration_files
        self.batch_idx = 0
        self.device_input = cuda.mem_alloc(
            batch_size * np.prod(input_shape) * 4
        )

    def get_batch(self, names):
        if self.batch_idx >= len(self.calibration_files):
            return None
        batch = []
        for i in range(self.batch_size):
            img = np.fromfile(
                self.calibration_files[self.batch_idx + i], 
                dtype=np.float32
            ).reshape(self.input_shape)
            batch.append(img)
        batch = np.stack(batch, axis=0)
        cuda.memcpy_htod(self.device_input, batch)
        self.batch_idx += self.batch_size
        return [int(self.device_input)]

    def get_batch_size(self):
        return self.batch_size

    def read_calibration_cache(self):
        return None

    def write_calibration_cache(self, cache):
        with open("calibration.cache", "wb") as f:
            f.write(cache)

校准集怎么选？我有个原则：覆盖产线所有常见缺陷类型。比如检测划痕，校准集里不能全是良品，得有划痕、压痕、脏污各占一些。数量上，500 张左右就够了，太多反而过拟合到校准集分布上。

量化方式	精度损失	推理加速	适用场景
FP32（基准）	无	1x	精度要求极高，算力充足
FP16	通常 < 0.5%	1.5x ~ 2x	大多数工业检测场景
INT8（熵校准）	1% ~ 3%	3x ~ 4x	缺陷明显、背景简单
INT8（分位数校准）	0.5% ~ 1.5%	3x ~ 4x	缺陷细小、需要保留细节

你想想看，INT8 加速 3 倍，但精度掉 2%，值不值？我个人的经验是：先跑 FP16，如果延迟达标就别折腾 INT8。INT8 的调试成本很高，而且不同批次的校准集可能导致结果波动。

4.4 整体流程一览

我把整个链路画了张图，方便你对照着操作：

这张图里我特意加了一条「反馈回路」。为什么？因为实际项目中，很少有一次就成功的。要么精度掉太多，要么延迟超标。这时候就得回到前面环节调整——换校准策略、改 opset、甚至重新训练时加入量化感知训练（QAT）。