第四章：模型准备——模型格式转换与量化对吞吐的影响

各位同学，今天我们来聊聊模型准备这个环节。说实话，很多做AI芯片压测的朋友，一上来就急着跑benchmark，结果模型格式不对、精度不对，折腾半天数据全是废的。我踩过这个坑，所以今天专门讲讲模型格式转换和量化对吞吐的影响。

4.1 模型格式转换：ONNX、TensorRT Engine、TFLite

模型格式转换，说白了就是把训练好的模型，变成你的芯片能“吃”的格式。不同的芯片，支持的格式不一样。比如NVIDIA的GPU，最擅长的是TensorRT Engine；移动端的芯片，比如高通、联发科，更习惯TFLite；而ONNX，是个中间格式，像个“翻译官”。

核心观点：模型格式转换不是简单的文件后缀改名，而是算子映射、内存布局、计算图优化的过程。转换得好，吞吐能翻倍；转换得不好，模型直接跑不起来。

4.1.1 ONNX：通用中间格式

ONNX（Open Neural Network Exchange）是微软和Meta联合推出的。它的好处是，你可以在PyTorch里训练，导出成ONNX，然后在TensorRT、OpenVINO、ONNX Runtime上跑。我个人习惯，先用ONNX做一次“格式验证”，看看模型能不能正常导出、推理结果对不对。

导出ONNX的代码很简单：

import torch
import torch.onnx

# 假设你有一个训练好的模型
model = torch.load('my_model.pth')
model.eval()

# 创建一个示例输入
dummy_input = torch.randn(1, 3, 224, 224)

# 导出ONNX
torch.onnx.export(
    model, 
    dummy_input, 
    'my_model.onnx',
    opset_version=11,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
)

嗯，这里要注意：opset_version 要选对。我遇到过，选了太高的版本，芯片的推理引擎不支持，结果模型导出来却跑不了。一般选11或13比较稳妥。

4.1.2 TensorRT Engine：GPU上的性能王者

TensorRT是NVIDIA的推理优化引擎。它能把ONNX模型进一步优化，比如层融合、精度校准、内存复用。我在项目中测试过，同样的模型，用TensorRT跑比直接用PyTorch推理，吞吐能提升3-5倍。

转换TensorRT Engine的流程：

  1. 准备好ONNX模型
  2. 用TensorRT的trtexec工具或Python API转换
  3. 指定精度（FP32、FP16、INT8）
  4. 生成.engine文件

命令行示例：

trtexec --onnx=my_model.onnx \
        --saveEngine=my_model.engine \
        --fp16 \
        --workspace=4096

为什么TensorRT能提升吞吐？因为它做了两件事：一是算子融合，把多个小算子合并成一个；二是内存优化，减少显存读写。你想想看，推理的时候，显存带宽往往是瓶颈，少读一次数据，速度就快一截。

我的经验：转换TensorRT Engine时，一定要检查--workspace大小。设得太小，转换会失败；设得太大，可能把显存撑爆。我一般设成4GB起步，根据模型大小调整。

4.1.3 TFLite：移动端和边缘设备的首选

TFLite是Google为移动端和嵌入式设备设计的。它支持量化、委派（Delegate）到GPU或NPU。如果你压测的是手机芯片或者边缘AI芯片，TFLite是绕不开的。

转换TFLite的代码：

import tensorflow as tf

# 加载Keras模型
model = tf.keras.models.load_model('my_model.h5')

# 转换器
converter = tf.lite.TFLiteConverter.from_keras_model(model)

# 启用量化
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float16]

# 转换
tflite_model = converter.convert()

# 保存
with open('my_model.tflite', 'wb') as f:
    f.write(tflite_model)

TFLite的量化支持很灵活。你可以选择FP16量化，也可以选择INT8量化。INT8量化后，模型大小能缩小到原来的1/4，推理速度也能提升2-3倍。但代价是精度会下降，具体下降多少，得看你的模型对量化是否敏感。

4.2 量化：INT8/FP16对吞吐的影响

量化，就是把模型里的浮点数（FP32）变成更小的数据类型（FP16或INT8）。为什么这么做？因为数据越小，计算越快，内存占用越少。说白了，就是用精度换速度。

我做过一个对比实验，结果如下：

精度类型 模型大小 推理延迟（ms） 吞吐（FPS） 精度损失（Top-1）
FP32 100% 12.5 80 0%
FP16 50% 7.2 139 0.1%
INT8 25% 4.1 244 0.8%

看到了吗？从FP32到FP16，吞吐提升了约74%；从FP32到INT8，吞吐提升了205%。但INT8的精度损失了0.8%。对于分类任务，0.8%的损失可以接受；但对于检测任务，尤其是小目标检测，这个损失可能就大了。

4.2.1 FP16量化：低风险，高回报

FP16量化，就是把32位浮点数变成16位。它的好处是，几乎不损失精度，但吞吐提升明显。为什么？因为现代GPU（比如NVIDIA的Turing、Ampere架构）都有专门的FP16计算单元，速度比FP32快一倍。

我在项目中遇到过，一个语义分割模型，FP32跑起来显存占用8GB，换成FP16后，显存降到4GB，吞吐从30FPS提升到55FPS。而且精度只掉了0.05%，肉眼根本看不出来。

注意：FP16量化不是所有算子都支持。有些算子（比如某些自定义的激活函数）在FP16下会溢出。我建议，量化后一定要做精度验证，跑一遍验证集，看看输出是否正常。

4.2.2 INT8量化：极致性能，但需要校准

INT8量化，是把浮点数变成8位整数。它的性能提升最大，但风险也最大。为什么？因为INT8的表示范围只有256个值，而FP32有40亿个值。量化过程中，信息丢失是必然的。

INT8量化的关键步骤是“校准”。你需要准备一个校准数据集（一般是几百到几千张图片），让模型在INT8下跑一遍，统计出每个激活值的分布，然后找到最优的量化参数（scale和zero_point）。

校准的代码示例（TensorRT）：

import tensorrt as trt

# 创建INT8校准器
class MyCalibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, calibration_files, batch_size):
        trt.IInt8EntropyCalibrator2.__init__(self)
        self.calibration_files = calibration_files
        self.batch_size = batch_size
        self.batch_idx = 0
        
    def get_batch(self, names):
        if self.batch_idx >= len(self.calibration_files):
            return None
        # 读取一批校准数据
        batch = self._load_batch(self.batch_idx)
        self.batch_idx += 1
        return [batch]
    
    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引擎
builder = trt.Builder(trt.Logger())
config = builder.create_builder_config()
config.set_flag(trt.BuilderFlag.INT8)
config.int8_calibrator = MyCalibrator(calibration_files, batch_size=32)

校准数据集的选择很重要。我吃过亏，用ImageNet的1000张图片做校准，结果模型在特定场景下精度掉了5%。后来换成目标场景的数据做校准，精度损失才降到1%以内。所以，校准数据集一定要贴近你的实际应用场景。

4.3 知识体系：模型准备的核心逻辑

为了让你更直观地理解，我画了一张流程图：

模型准备核心流程 训练好的模型 模型格式转换 ONNX TensorRT Engine TFLite 量化（FP16 / INT8） 吞吐压测

这张图很清晰：训练好的模型，先做格式转换（ONNX、TensorRT Engine、TFLite），再做量化（FP16或INT8），最后才能进行吞吐压测。每一步都有坑，每一步都影响最终结果。

4.4 避坑指南

最后，分享几个我踩过的坑：

  • 我曾经直接拿FP32模型跑压测，结果显存爆了。后来才发现，模型没做量化，显存占用是INT8的4倍。所以，压测前一定要确认模型格式和精度。
  • 我曾经用TensorRT转换ONNX模型，结果报错“unsupported op”。查了半天，发现是ONNX里有个自定义算子，TensorRT不认识。解决办法是，要么替换成标准算子，要么写插件。
  • 我曾经做INT8量化，校准数据集选错了，结果模型精度掉了10%。后来换成目标场景的数据，精度才恢复正常。所以，校准数据集一定要有代表性。

好了，模型准备这部分就讲到这里。记住，格式转换和量化是吞吐压测的“前戏”，做得好，后面事半功倍；做得不好，数据全是假的。下节课，我们聊聊压测工具的选择和使用。

想要更多资料，联系可可

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