4. 模型导出与优化：ONNX导出、FP16/INT8量化、TensorRT引擎构建、精度验证

好，咱们进入实战环节最核心的一步——模型导出与优化。这一步说白了，就是把训练好的PyTorch模型，变成能在车规级芯片上真正跑起来的“子弹”。我个人习惯把这步叫做“炼金术”，因为处理得好，模型又快又准；处理不好，直接废掉。

4.1 ONNX导出：打通PyTorch到TensorRT的桥梁

ONNX（Open Neural Network Exchange）就像一个通用语言。PyTorch训练好的模型，得先翻译成ONNX格式，TensorRT才能读懂。我在项目中遇到过最头疼的问题，就是导出时算子不兼容。

⚠️ 避坑指南： 我曾经因为一个torch.where操作，导致ONNX导出后多出三个Reshape节点，推理速度直接翻倍。后来我改用torch.clamp + 掩码乘法替代，问题才解决。

导出代码其实不复杂，但有几个关键点要注意：

import torch
import torch.onnx

def export_onnx(model, dummy_input, save_path):
    model.eval()
    # 设置动态轴，这对ADAS场景很重要
    # 因为实际部署时，输入图像的宽高可能变化
    dynamic_axes = {
        'input': {0: 'batch_size', 2: 'height', 3: 'width'},
        'output': {0: 'batch_size'}
    }
    
    torch.onnx.export(
        model,
        dummy_input,
        save_path,
        opset_version=11,  # 我建议用11或12，兼容性好
        do_constant_folding=True,
        input_names=['input'],
        output_names=['output'],
        dynamic_axes=dynamic_axes
    )
    print(f"ONNX模型已保存至: {save_path}")

嗯，这里要注意opset_version的选择。版本太低，有些新算子不支持；版本太高，老版本的TensorRT可能不认。我个人习惯用opset 11，兼容性最好。

4.2 FP16/INT8量化：精度与速度的博弈

量化，说白了就是把模型从FP32的“高精度慢跑”，变成FP16或INT8的“低精度快跑”。在ADAS里，速度就是生命。你想想看，一个行人检测模型，FP32跑30ms，INT8跑8ms，这差距意味着什么？

4.2.1 FP16量化

FP16量化最简单，几乎无痛。大部分GPU都原生支持FP16计算。我建议你先把模型转成FP16试试，如果精度损失在1%以内，那就直接用。

# TensorRT中启用FP16
import tensorrt as trt

def build_fp16_engine(onnx_path, engine_path):
    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)
    
    with open(onnx_path, 'rb') as f:
        parser.parse(f.read())
    
    config = builder.create_builder_config()
    config.set_flag(trt.BuilderFlag.FP16)  # 启用FP16
    config.max_workspace_size = 1 << 30  # 1GB
    
    engine = builder.build_engine(network, config)
    with open(engine_path, 'wb') as f:
        f.write(engine.serialize())
    print(f"FP16引擎已保存至: {engine_path}")
💡 小技巧： 如果FP16精度损失较大，可以尝试混合精度——让敏感层保持FP32，其他层用FP16。TensorRT的set_precision接口可以做到。

4.2.2 INT8量化

INT8量化就复杂多了。它需要校准数据集，而且对模型结构敏感。我在项目中遇到过，同一个YOLOv8模型，在白天场景下INT8量化后mAP只掉0.5%，但到了夜间场景直接掉3%。

为什么会这样？因为校准数据集没有覆盖夜间场景。所以，校准数据一定要有代表性。

# INT8量化需要校准器
class Calibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, calibration_files, batch_size):
        trt.IInt8EntropyCalibrator2.__init__(self)
        self.calibration_files = calibration_files
        self.batch_size = batch_size
        self.index = 0
        # 分配GPU内存
        self.device_input = cuda.mem_alloc(1 * 3 * 640 * 640 * 4)
    
    def get_batch(self, names):
        if self.index < len(self.calibration_files):
            # 读取并预处理图像
            batch = self._load_batch()
            cuda.memcpy_htod(self.device_input, batch)
            self.index += 1
            return [int(self.device_input)]
        else:
            return None
    
    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)

# 构建INT8引擎
config.set_flag(trt.BuilderFlag.INT8)
config.int8_calibrator = Calibrator(calibration_files, batch_size=8)
🔑 关键点： INT8量化后一定要做精度验证。我见过太多人量化完直接部署，结果在高速上把卡车检测成行人——这可不是闹着玩的。

4.3 TensorRT引擎构建：从ONNX到可执行引擎

TensorRT引擎构建，就是把ONNX模型编译成针对特定GPU优化的可执行文件。这个过程会做很多优化：层融合、内存复用、内核自动调优等。

我个人习惯把引擎构建和推理分开。引擎构建一次，推理多次。毕竟构建引擎很耗时，在车上不可能每次启动都重新构建。

def build_engine(onnx_path, engine_path, precision='fp16'):
    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(onnx_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.max_workspace_size = 1 << 30  # 1GB
    
    if precision == 'fp16':
        config.set_flag(trt.BuilderFlag.FP16)
    elif precision == 'int8':
        config.set_flag(trt.BuilderFlag.INT8)
        config.int8_calibrator = Calibrator(calibration_files, batch_size=8)
    
    # 设置最大batch size
    profile = builder.create_optimization_profile()
    profile.set_shape('input', (1, 3, 640, 640), (4, 3, 640, 640), (8, 3, 640, 640))
    config.add_optimization_profile(profile)
    
    engine = builder.build_engine(network, config)
    with open(engine_path, 'wb') as f:
        f.write(engine.serialize())
    return engine
⚠️ 注意： 不同GPU架构（如Turing、Ampere、Hopper）生成的引擎不通用。你在开发机上用RTX 4090生成的引擎，不能直接部署到车上的Orin上。必须重新构建。

4.4 精度验证：别让优化毁了模型

精度验证是最后一道防线。我见过太多人优化完模型，跑起来飞快，但检测结果一塌糊涂。所以，一定要量化评估。

我建议至少做三件事：

  1. 逐层精度对比：对比FP32和优化后模型每一层的输出差异
  2. 整体mAP评估：在验证集上计算mAP，确保下降在可接受范围内
  3. 可视化检查：随机抽几张图，肉眼对比检测结果
def validate_precision(original_model, optimized_model, dataloader):
    original_mAP = evaluate(original_model, dataloader)
    optimized_mAP = evaluate(optimized_model, dataloader)
    
    drop = original_mAP - optimized_mAP
    print(f"原始模型mAP: {original_mAP:.4f}")
    print(f"优化后mAP: {optimized_mAP:.4f}")
    print(f"精度下降: {drop:.4f}")
    
    if drop > 0.05:  # 下降超过5%
        print("⚠️ 警告：精度下降过大，建议检查量化策略")
    else:
        print("✅ 精度验证通过")
💡 经验之谈： 我曾经遇到一个案例，INT8量化后mAP只掉了2%，但实际路测时，模型对远距离行人的召回率从85%掉到了60%。后来发现是校准数据里缺少远距离样本。所以，验证集一定要覆盖所有典型场景。

4.5 总结：优化流程一览

步骤 操作 关键点 耗时
1 ONNX导出 算子兼容性、动态轴设置 几分钟
2 FP16量化 几乎无痛，直接启用 几分钟
3 INT8量化 校准数据集要有代表性 几十分钟
4 TensorRT引擎构建 GPU架构匹配、优化profile 几十分钟
5 精度验证 逐层对比、mAP评估、可视化 几小时

嗯，整个流程走下来，快的话半天，慢的话一两天。但别嫌麻烦，这一步做好了，后面部署就稳了。我见过太多项目，前面训练花了一个月，优化只花了一天，结果上线就出问题。说白了，优化这一步，值得你花时间。

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