第4章：模型转换与优化：从PyTorch/TensorFlow导出ONNX、FP32转FP16/INT8量化、模型剪枝与蒸馏基础

说实话，模型训练出来只是第一步。真正要把模型塞进手机、摄像头或者边缘盒子，你会发现——嗯，事情没那么简单。

我最早做嵌入式AI的时候，踩过一个坑。在服务器上跑得飞快的ResNet，转到树莓派上，推理一次要3秒。这谁受得了？后来才明白，模型转换和优化，才是端侧部署的真正门槛。

这一章，我们就来聊聊这个核心话题。我会把我在项目中积累的经验，尽量说透。

4.1 从PyTorch/TensorFlow导出ONNX

ONNX是什么？说白了，它是一个中间格式。就像你写了一份中文文档，要翻译成英文、日文、法文。ONNX就是那个“中间语言”。

我个人习惯，不管用什么框架训练，最后都转成ONNX。为什么？因为几乎所有端侧推理引擎都支持它。你想想看，如果每个框架都写一套转换代码，那得多累。

4.1.1 PyTorch导出ONNX

PyTorch导出ONNX其实很简单。核心就一个函数：torch.onnx.export()。但有几个坑，我当年可是一个个踩过来的。

import torch
import torchvision.models as models

# 加载一个预训练模型
model = models.resnet18(pretrained=True)
model.eval()

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

# 导出 ONNX
torch.onnx.export(
    model,               # 模型
    dummy_input,         # 输入张量
    "resnet18.onnx",     # 输出文件名
    export_params=True,  # 是否导出参数
    opset_version=11,    # ONNX算子集版本
    do_constant_folding=True,  # 常量折叠优化
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

这里有个关键点：opset_version。我建议用11或以上。版本太低，有些算子不支持。版本太高，有些推理引擎又不认识。11是个比较稳妥的选择。

注意： 导出时一定要调用 model.eval()。我见过有人忘了这步，结果导出的模型里还带着Dropout和BatchNorm的训练行为，推理结果完全不对。

4.1.2 TensorFlow导出ONNX

TensorFlow这边，我一般用 tf2onnx 这个工具。它支持SavedModel和Keras模型两种格式。

import tensorflow as tf
import tf2onnx

# 加载Keras模型
model = tf.keras.applications.MobileNetV2(weights='imagenet')

# 指定输入签名
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)

# 转换
output_path = model.output_names[0]
model_proto, _ = tf2onnx.convert.from_keras(
    model, 
    input_signature=spec,
    opset=13,
    output_path="mobilenetv2.onnx"
)

我曾经遇到过一个情况：TensorFlow模型里用了自定义层，结果转ONNX时报错。后来发现，需要先注册自定义算子。这个坑，建议大家提前检查模型里有没有非标准操作。

4.1.3 ONNX模型验证

导出之后，一定要验证。别问我为什么强调这个——我吃过亏。

import onnx
import onnxruntime as ort

# 检查模型结构
onnx_model = onnx.load("resnet18.onnx")
onnx.checker.check_model(onnx_model)
print("模型结构检查通过")

# 用ONNX Runtime跑一下推理
session = ort.InferenceSession("resnet18.onnx")
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

import numpy as np
dummy = np.random.randn(1, 3, 224, 224).astype(np.float32)
result = session.run([output_name], {input_name: dummy})
print("推理结果形状:", result[0].shape)
小技巧： 可以用 onnx.shape_inference.infer_shapes() 来检查张量形状是否合理。如果某个中间层的形状是负数，说明模型可能有问题。

4.2 FP32转FP16/INT8量化

模型转成ONNX之后，下一步就是量化。为什么要量化？说白了，就是让模型变小、跑得更快。

FP32的模型，一个参数占4字节。FP16只占2字节，INT8只占1字节。你想想看，模型大小直接减半甚至减到四分之一。在嵌入式设备上，这差别太大了。

4.2.1 FP16量化

FP16量化是最简单的。几乎不需要额外数据，直接转就行。精度损失也很小，通常不到1%。

import onnx
from onnxconverter_common import float16

# 加载FP32模型
model_fp32 = onnx.load("resnet18.onnx")

# 转FP16
model_fp16 = float16.convert_float_to_float16(model_fp32)

# 保存
onnx.save(model_fp16, "resnet18_fp16.onnx")

我个人习惯，对于视觉模型，FP16基本够用。但如果是语音或某些对精度极其敏感的模型，建议先做对比测试。

经验之谈： 我在一个车牌识别项目里，FP32转FP16后，准确率从98.7%降到了98.5%。几乎可以忽略。但模型推理速度提升了40%。这笔账，划算。

4.2.2 INT8量化

INT8量化就复杂多了。它需要校准数据，而且精度损失可能更大。但收益也更大——模型大小只有FP32的四分之一。

INT8量化有两种方式：

  • 训练后量化（PTQ）：训练完模型后，用少量校准数据做量化。简单快速，但精度可能下降。
  • 量化感知训练（QAT）：在训练过程中模拟量化效果。精度更高，但需要重新训练。

这里我重点讲PTQ，因为它在实际项目里用得最多。

import onnx
from onnxruntime.quantization import quantize_dynamic, QuantType

# 动态量化
quantized_model = quantize_dynamic(
    "resnet18.onnx",
    "resnet18_int8.onnx",
    weight_type=QuantType.QInt8
)

等等，上面这个只是权重量化。如果你想做全量化（权重+激活），需要准备校准数据：

from onnxruntime.quantization import quantize_static, CalibrationMethod
from onnxruntime.quantization import CalibrationDataReader

class MyCalibrationDataReader(CalibrationDataReader):
    def __init__(self):
        self.data = [np.random.randn(1, 3, 224, 224).astype(np.float32) for _ in range(100)]
        self.index = 0
    
    def get_next(self):
        if self.index < len(self.data):
            result = {"input": self.data[self.index]}
            self.index += 1
            return result
        return None

# 静态量化
quantize_static(
    "resnet18.onnx",
    "resnet18_int8_static.onnx",
    calibration_data_reader=MyCalibrationDataReader(),
    calibration_method=CalibrationMethod.MinMax
)
注意： INT8量化后，一定要做精度验证。我曾经在一个项目里，量化后模型准确率从95%掉到了82%。后来发现是校准数据分布和真实数据不一致。校准数据一定要有代表性。

4.3 模型剪枝与蒸馏基础

量化和剪枝、蒸馏，这三兄弟经常一起用。量化是“压缩数值”，剪枝是“砍掉冗余”，蒸馏是“以大教小”。

4.3.1 模型剪枝

剪枝，说白了就是把不重要的连接或通道删掉。我最早接触剪枝时，觉得这玩意儿很玄学。后来发现，其实原理很简单——那些权重接近0的参数，删了也没太大影响。

剪枝分两种：

  • 非结构化剪枝：把单个权重置为0。稀疏度高，但硬件加速难。
  • 结构化剪枝：直接删掉整个通道或卷积核。硬件友好，但精度损失可能更大。

PyTorch里做剪枝很方便：

import torch
import torch.nn.utils.prune as prune

# 对卷积层做L1非结构化剪枝
model = models.resnet18(pretrained=True)
prune.l1_unstructured(
    model.conv1,
    name='weight',
    amount=0.3  # 剪掉30%的权重
)

# 移除剪枝掩码，永久生效
prune.remove(model.conv1, 'weight')

结构化剪枝稍微复杂一点：

# 对卷积层做L1结构化剪枝（按通道）
prune.ln_structured(
    model.layer1[0].conv1,
    name='weight',
    amount=0.2,  # 剪掉20%的通道
    n=2,         # L2范数
    dim=0        # 按输出通道剪
)
我的建议： 剪枝比例不要超过50%。我试过剪掉70%，模型直接废了。一般20%-30%是比较安全的范围。剪完后最好做一下微调（fine-tune），能挽回不少精度。

4.3.2 知识蒸馏

知识蒸馏，就是用一个大的教师模型，去教一个小学生模型。教师模型输出的是“软标签”——不仅仅是类别，还有各类别的概率分布。

蒸馏的核心公式：

# 蒸馏损失 = 蒸馏损失权重 * KL散度(学生输出, 教师输出) + (1-蒸馏权重) * 交叉熵(学生输出, 真实标签)

def distillation_loss(student_logits, teacher_logits, labels, T=4.0, alpha=0.7):
    """
    T: 温度参数，越大软标签越平滑
    alpha: 蒸馏损失权重
    """
    # 软标签损失
    soft_loss = nn.KLDivLoss()(
        F.log_softmax(student_logits / T, dim=1),
        F.softmax(teacher_logits / T, dim=1)
    ) * (T * T)
    
    # 硬标签损失
    hard_loss = nn.CrossEntropyLoss()(student_logits, labels)
    
    return alpha * soft_loss + (1 - alpha) * hard_loss

我做过一个实验：用ResNet-50当教师，教一个MobileNet-V2学生。蒸馏后，MobileNet的准确率从72%提升到了76%。虽然还是比不上教师（78%），但模型大小只有教师的十分之一。

关键参数： 温度T一般设4-8。太高了软标签太平滑，学不到东西。太低了又跟硬标签没区别。我习惯从T=4开始试，效果不好再调。

4.4 综合优化流程

在实际项目中，我一般按这个流程走：

  1. 训练FP32模型，确保基线精度达标
  2. 导出ONNX，验证结构正确
  3. 尝试FP16量化，如果精度损失可接受，直接用
  4. 如果FP16不够，做剪枝+微调，再转INT8
  5. 如果INT8精度不够，用蒸馏训练一个小模型，再量化

这个流程，我用了很多项目，基本都能搞定。当然，具体问题要具体分析。比如有些模型对精度极其敏感，那就只能牺牲一些速度了。

最后说一句：模型优化没有银弹。量化、剪枝、蒸馏，每个方法都有适用场景。关键是多试、多测、多对比。我在项目中养成的习惯是——每次优化后，都跑一遍完整的验证集。别偷懒，这个步骤省不得。

下一章，我们会聊推理引擎的运行时架构。到时候，我会讲讲内存管理和算子调度那些事。嗯，那又是另一番天地了。

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