模型剪枝与量化联合优化：从结构压缩到精度补偿的边缘 AI 工程链路

一、单一优化的瓶颈：剪枝后精度崩了，量化后跑不动

边缘 AI 部署中，模型压缩是绕不开的环节。常用的两种压缩手段——剪枝（Pruning）和量化（Quantization）——各自有效，但单独使用时都有明显瓶颈。剪枝可以减少 50-70% 的参数量，但剪枝后的稀疏模型如果不经过专门优化，在大多数边缘推理引擎上反而更慢（稀疏矩阵运算效率低于稠密矩阵）。量化可以将模型体积压缩 4 倍（FP32 → INT8），但对某些层（如注意力机制、残差连接）量化后精度损失严重。

联合优化的思路是：先剪枝移除冗余参数，再量化压缩数值精度，最后通过知识蒸馏（Knowledge Distillation）补偿精度损失。三步串联，在延迟、体积和精度三个维度同时达标。这不是简单的"1+1"，而是需要精心编排的工程链路。

二、剪枝-量化-蒸馏联合优化链路

联合优化的关键在于顺序和粒度控制。先剪枝再量化，是因为剪枝改变了模型结构，量化的校准范围需要基于剪枝后的模型重新计算。知识蒸馏放在最后，是因为剪枝和量化都会引入精度损失，蒸馏可以一次性补偿两种损失。

flowchart TD
    A[原始 FP32 模型] --> B[结构化剪枝]
    B --> C{稀疏率评估}
    C -->|精度达标| D[微调恢复]
    C -->|精度不达标| E[降低稀疏率]
    E --> B

    D --> F[INT8 静态量化]
    F --> G{量化精度评估}
    G -->|精度达标| H[知识蒸馏]
    G -->|精度不达标| I[混合精度量化]
    I --> H

    H --> J[学生模型训练]
    J --> K{最终精度评估}
    K -->|达标| L[导出部署模型]
    K -->|不达标| M[调整蒸馏温度/损失权重]
    M --> J

    style B fill:#fbb,stroke:#333
    style F fill:#bbf,stroke:#333
    style H fill:#bfb,stroke:#333

2.1 结构化剪枝 vs 非结构化剪枝

• 非结构化剪枝：逐参数置零，稀疏度高但需要专门引擎支持
• 结构化剪枝：按通道/层整块移除，稀疏度略低但无需特殊引擎，推理速度直接提升

边缘部署中，结构化剪枝是首选，因为 ONNX Runtime、TFLite 等引擎对结构化稀疏有原生支持。

2.2 量化敏感度分析

并非所有层都适合 INT8 量化。第一层卷积和最后一层全连接通常对量化最敏感。通过逐层量化敏感度分析，可以识别出需要保持高精度的层，采用混合精度策略。

三、生产级代码实现

3.1 结构化通道剪枝

# channel_pruning.py
# 基于 BN 权重的结构化通道剪枝
import torch
import torch.nn as nn
import numpy as np
from typing import list


def compute_bn_importance(model: nn.Module) -> dict[str, np.ndarray]:
    """基于 BN 层 gamma 权重计算通道重要性"""
    importance = {}
    for name, module in model.named_modules():
        if isinstance(module, nn.BatchNorm2d):
            # gamma 绝对值越大，该通道越重要
            gamma = module.weight.data.abs().cpu().numpy()
            importance[name] = gamma
    return importance


def get_pruning_mask(
    importance: np.ndarray,
    prune_ratio: float
) -> np.ndarray:
    """根据重要性分数生成剪枝掩码"""
    threshold = np.sort(importance)[
        int(len(importance) * prune_ratio)
    ]
    return importance > threshold


def prune_conv_bn(
    conv: nn.Conv2d,
    bn: nn.BatchNorm2d,
    mask: np.ndarray
):
    """对 Conv + BN 结构执行通道剪枝"""
    # 保留的通道索引
    keep_indices = np.where(mask)[0]

    # 剪枝 Conv 输出通道
    conv.weight = nn.Parameter(
        conv.weight.data[keep_indices]
    )
    if conv.bias is not None:
        conv.bias = nn.Parameter(
            conv.bias.data[keep_indices]
        )
    conv.out_channels = len(keep_indices)

    # 剪枝 BN 参数
    bn.weight = nn.Parameter(bn.weight.data[keep_indices])
    bn.bias = nn.Parameter(bn.bias.data[keep_indices])
    bn.running_mean = bn.running_mean[keep_indices]
    bn.running_var = bn.running_var[keep_indices]
    bn.num_features = len(keep_indices)

    return keep_indices


def prune_model(
    model: nn.Module,
    prune_ratio: float = 0.3
) -> dict[str, list[int]]:
    """对整个模型执行结构化通道剪枝"""
    importance = compute_bn_importance(model)
    pruned_indices = {}

    for name, module in model.named_modules():
        if isinstance(module, nn.BatchNorm2d) and name in importance:
            mask = get_pruning_mask(importance[name], prune_ratio)
            # 找到对应的 Conv 层
            parent_name = ".".join(name.split(".")[:-1])
            conv = dict(model.named_modules()).get(
                f"{parent_name}.conv"
            ) or dict(model.named_modules()).get(
                f"{parent_name}.0"
            )
            if conv and isinstance(conv, nn.Conv2d):
                keep = prune_conv_bn(conv, module, mask)
                pruned_indices[name] = keep.tolist()

    return pruned_indices

3.2 逐层量化敏感度分析

# quantization_sensitivity.py
# 逐层量化敏感度分析：找出不适合 INT8 量化的层
import numpy as np
import onnxruntime as ort
from onnxruntime.quantization import (
    quantize_static,
    CalibrationDataReader,
    QuantType,
    QuantFormat
)


class SensitivityAnalyzer:
    """逐层量化敏感度分析器"""

    def __init__(
        self,
        model_path: str,
        calibration_data: np.ndarray
    ):
        self.model_path = model_path
        self.calibration_data = calibration_data
        # 基线精度（FP32）
        self.baseline_output = self._run_inference(model_path)

    def _run_inference(self, model_path: str) -> np.ndarray:
        session = ort.InferenceSession(
            model_path, providers=["CPUExecutionProvider"]
        )
        input_name = session.get_inputs()[0].name
        return session.run(
            None,
            {input_name: self.calibration_data.astype(np.float32)}
        )[0]

    def analyze_layer(
        self,
        layer_name: str,
        quantize_this_layer: bool = True
    ) -> float:
        """分析单个层的量化敏感度"""
        # 量化时排除/包含目标层
        nodes_to_exclude = [] if quantize_this_layer else [layer_name]

        output_path = f"/tmp/sensitivity_{layer_name}.onnx"
        quantize_static(
            self.model_path,
            output_path,
            CalibrationDataReader(self.calibration_data),
            quant_format=QuantFormat.QDQ,
            weight_type=QuantType.QInt8,
            nodes_to_exclude=nodes_to_exclude
        )

        quantized_output = self._run_inference(output_path)
        # 计算输出差异（MSE）
        mse = np.mean((self.baseline_output - quantized_output) ** 2)
        return float(mse)

    def full_analysis(self, layer_names: list[str]) -> dict[str, float]:
        """对所有层执行敏感度分析"""
        results = {}
        for name in layer_names:
            mse = self.analyze_layer(name, quantize_this_layer=True)
            results[name] = mse
            print(f"  {name}: MSE = {mse:.6f}")

        # 按敏感度排序
        sorted_results = dict(
            sorted(results.items(), key=lambda x: -x[1])
        )
        return sorted_results

3.3 知识蒸馏精度补偿

# knowledge_distillation.py
# 剪枝+量化后的知识蒸馏精度补偿
import torch
import torch.nn as nn
import torch.nn.functional as F


class DistillationTrainer:
    """知识蒸馏训练器"""

    def __init__(
        self,
        teacher: nn.Module,
        student: nn.Module,
        temperature: float = 4.0,
        alpha: float = 0.7,
        learning_rate: float = 1e-4
    ):
        self.teacher = teacher
        self.student = student
        self.temperature = temperature
        self.alpha = alpha  # 蒸馏损失权重

        # 教师模型冻结，不参与梯度更新
        for p in self.teacher.parameters():
            p.requires_grad = False
        self.teacher.eval()

        self.optimizer = torch.optim.AdamW(
            self.student.parameters(), lr=learning_rate
        )

    def distillation_loss(
        self,
        student_logits: torch.Tensor,
        teacher_logits: torch.Tensor,
        labels: torch.Tensor
    ) -> torch.Tensor:
        """蒸馏损失 = alpha * KL散度 + (1-alpha) * 交叉熵"""
        # 软标签蒸馏损失
        soft_student = F.log_softmax(
            student_logits / self.temperature, dim=-1
        )
        soft_teacher = F.softmax(
            teacher_logits / self.temperature, dim=-1
        )
        kl_loss = F.kl_div(
            soft_student, soft_teacher, reduction="batchmean"
        ) * (self.temperature ** 2)

        # 硬标签分类损失
        ce_loss = F.cross_entropy(student_logits, labels)

        return self.alpha * kl_loss + (1 - self.alpha) * ce_loss

    def train_step(
        self,
        inputs: torch.Tensor,
        labels: torch.Tensor
    ) -> float:
        """单步蒸馏训练"""
        self.student.train()

        with torch.no_grad():
            teacher_logits = self.teacher(inputs)

        student_logits = self.student(inputs)
        loss = self.distillation_loss(
            student_logits, teacher_logits, labels
        )

        self.optimizer.zero_grad()
        loss.backward()
        # 梯度裁剪：防止蒸馏初期梯度爆炸
        torch.nn.utils.clip_grad_norm_(
            self.student.parameters(), max_norm=1.0
        )
        self.optimizer.step()

        return loss.item()

3.4 联合优化流水线

# joint_optimization_pipeline.py
# 剪枝 → 量化 → 蒸馏 联合优化流水线
import torch
from channel_pruning import prune_model
from quantization_sensitivity import SensitivityAnalyzer
from knowledge_distillation import DistillationTrainer


class JointOptimizationPipeline:
    """联合优化流水线"""

    def __init__(
        self,
        model: torch.nn.Module,
        train_loader: torch.utils.data.DataLoader,
        val_loader: torch.utils.data.DataLoader,
        prune_ratio: float = 0.3,
        target_accuracy: float = 0.95
    ):
        self.original_model = model
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.prune_ratio = prune_ratio
        self.target_accuracy = target_accuracy

    def run(self):
        """执行完整的联合优化流程"""
        model = self.original_model

        # Step 1: 结构化剪枝
        print(f"=== Step 1: 结构化剪枝 (ratio={self.prune_ratio}) ===")
        pruned_indices = prune_model(model, self.prune_ratio)
        pruned_params = sum(
            p.numel() for p in model.parameters()
        )
        original_params = sum(
            p.numel() for p in self.original_model.parameters()
        )
        print(
            f"参数量: {original_params} → {pruned_params} "
            f"(压缩率 {pruned_params/original_params:.1%})"
        )

        # Step 1.5: 剪枝后微调恢复精度
        print("=== Step 1.5: 剪枝后微调 ===")
        self._finetune(model, epochs=5)

        # Step 2: 量化敏感度分析 + INT8 量化
        print("=== Step 2: 量化敏感度分析 ===")
        # 导出 ONNX 用于量化分析
        self._export_onnx(model, "/tmp/pruned_model.onnx")
        analyzer = SensitivityAnalyzer(
            "/tmp/pruned_model.onnx",
            calibration_data=self._get_calibration_data()
        )
        sensitive_layers = analyzer.full_analysis(
            self._get_conv_layer_names(model)
        )

        # Step 3: 知识蒸馏补偿精度
        print("=== Step 3: 知识蒸馏 ===")
        distiller = DistillationTrainer(
            teacher=self.original_model,
            student=model,
            temperature=4.0,
            alpha=0.7
        )
        self._distill(distiller, epochs=10)

        # 最终评估
        accuracy = self._evaluate(model)
        print(f"最终精度: {accuracy:.4f} (目标: {self.target_accuracy})")

        return model

    def _finetune(self, model, epochs: int):
        optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
        criterion = torch.nn.CrossEntropyLoss()
        for epoch in range(epochs):
            model.train()
            for inputs, labels in self.train_loader:
                optimizer.zero_grad()
                loss = criterion(model(inputs), labels)
                loss.backward()
                optimizer.step()

    def _distill(self, distiller, epochs: int):
        for epoch in range(epochs):
            total_loss = 0
            for inputs, labels in self.train_loader:
                loss = distiller.train_step(inputs, labels)
                total_loss += loss
            avg_loss = total_loss / len(self.train_loader)
            print(f"  Epoch {epoch+1}: distill_loss={avg_loss:.4f}")

    def _evaluate(self, model) -> float:
        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for inputs, labels in self.val_loader:
                outputs = model(inputs)
                _, predicted = outputs.max(1)
                total += labels.size(0)
                correct += predicted.eq(labels).sum().item()
        return correct / total

四、联合优化的工程代价：训练成本、流水线复杂度与精度天花板

联合优化不是免费的午餐，以下 Trade-offs 需要在工程决策中权衡：

训练成本倍增。剪枝微调 + 量化校准 + 蒸馏训练，整个流水线的训练成本是单纯微调的 3-5 倍。在资源受限的团队中，需要评估投入产出比——如果模型本身已经足够小（如 MobileNetV3），可能只需要量化就够了，不需要走完整的联合优化流程。

流水线复杂度。三步串联意味着三处可能失败的环节。剪枝后精度崩了需要回退稀疏率，量化后精度不达标需要切换混合精度，蒸馏后精度仍不够需要调整温度和损失权重。每一步都需要人工判断和调参，自动化程度有限。建议建立标准化的评估检查点：剪枝后精度不低于基线的 95%，量化后精度不低于剪枝后模型的 97%，蒸馏后精度不低于基线的 99%。

精度天花板。联合优化能补偿大部分精度损失，但无法完全恢复。经验上，剪枝 30% + INT8 量化的联合方案，最终精度通常比原始 FP32 模型低 1-2%。如果业务要求精度损失不超过 0.5%，可能需要降低剪枝比例或使用更高精度的量化方案（如 FP16 代替 INT8）。

部署兼容性。混合精度量化（部分层 INT8，部分层 FP16/FP32）在 ONNX Runtime 上支持良好，但在某些 NPU 上可能不支持混合精度推理，所有层必须统一为 INT8。部署前必须确认目标硬件的量化格式支持情况。

五、总结

剪枝-量化-蒸馏联合优化的核心价值在于，通过三步串联在延迟、体积和精度三个维度同时达标，而非单一维度的优化。落地要点如下：

1. 结构化剪枝优先：使用 BN 权重评估通道重要性，结构化剪枝直接减少计算量，无需稀疏推理引擎
2. 敏感度分析驱动量化：逐层分析量化敏感度，对敏感层保持高精度，避免"一刀切"量化导致精度崩塌
3. 蒸馏补偿精度：剪枝和量化引入的精度损失，通过知识蒸馏一次性补偿，温度参数和损失权重需要调优
4. 检查点评估：每步完成后评估精度，设定明确的回退阈值，避免在不可逆的精度损失后继续优化
5. 按需组合：小模型只需量化，大模型才需要完整的三步联合优化，避免过度工程化