02 模型压缩

学习目标

理解模型压缩的概念、目的和主要方法
掌握量化技术的原理，包括线性量化、对称/非对称量化
了解训练后量化（PTQ）和量化感知训练（QAT）的流程与方法
掌握模型剪枝的基本原理，包括结构化剪枝和非结构化剪枝
理解知识蒸馏的核心思想和经典算法

2.1 模型压缩概述

2.1.1 什么是模型压缩

随着深度学习技术的快速发展，神经网络的规模变得越来越大。从早期的LeNet、AlexNet到如今的GPT、BERT等大模型，网络参数从数百万增长到数千亿甚至更多。这种规模的增长带来的好处是模型性能的显著提升，但同时也带来了严峻的部署挑战：

存储成本高昂：大模型的参数文件可能达到数GB甚至上百GB，难以存储和传输
计算资源需求大：大模型的推理需要大量的浮点运算，推理延迟高
内存占用惊人：大模型加载到内存中需要占用大量存储空间，边缘设备难以承载
能耗问题突出：移动设备依赖电池供电，高计算量的模型会快速消耗电量

模型压缩（Model Compression）是一类技术的统称，旨在保持模型性能的前提下，减小模型的存储体积、降低计算复杂度和减少内存占用。模型压缩不是简单地删除模型，而是通过各种技术手段重新表示模型信息，在"表示能力"和"压缩率"之间找到最优平衡。

模型压缩的目标可以用数学语言描述：给定一个训练好的原始模型\(M\)和压缩后的模型\(M'\)，在满足压缩率约束（如参数量减少4倍）的前提下，最小化性能损失（如精度下降不超过1%）。

2.1.2 模型压缩的主要方法

模型压缩技术可以分为四大类别：量化（Quantization）、剪枝（Pruning）、知识蒸馏（Knowledge Distillation）和低秩分解（Low-Rank Factorization）。

量化通过减少模型参数和计算的数值精度来减小模型体积。典型的量化方法将FP32浮点数转换为INT8整数表示，可以将模型体积减少4倍，同时加速推理计算（因为整数运算比浮点运算更快更省电）。

剪枝通过删除模型中"不重要"的参数来减小模型规模。这里的"不重要"可以理解为对模型输出影响较小的权重或结构。剪枝可以是粗粒度的（删除整个神经元或卷积核）或细粒度的（删除单个权重连接）。

知识蒸馏通过让"学生模型"学习"教师模型"的知识来实现模型压缩。教师模型通常是大而精确的模型，学生模型是小而高效的模型。蒸馏过程中，学生模型不仅学习真实标签，还学习教师模型的输出分布（软标签）。

低秩分解通过将大矩阵分解为多个小矩阵的乘积来减少参数量。例如，将一个\(1000 \times 1000\)的矩阵分解为\(1000 \times 100 \times 100 \times 1000\)的形式，参数量从100万减少到20万。

这四类方法各有优劣，实际应用中经常组合使用以获得更好的压缩效果。

2.1.3 模型压缩的评估指标

评估模型压缩的效果需要综合考虑多个维度：

压缩率（Compression Ratio）

压缩率是最直观的指标，定义为原始模型大小与压缩后模型大小的比值。如果原始模型是100MB，压缩后是25MB，则压缩率为4倍。

压缩率可以分别从参数量、模型文件大小、计算量等角度计算：

参数压缩率：基于参数量计算的压缩效果
体积压缩率：基于模型文件大小计算的压缩效果
计算压缩率：基于FLOPs（浮点运算次数）计算的压缩效果

精度损失（Accuracy Loss）

压缩后的模型在目标任务上的精度与原始模型的精度之差。不同的应用场景对精度损失有不同的容忍度：

图像分类：通常允许1-2%的精度损失
目标检测：通常允许2-3%的精度损失
自动驾驶：通常要求精度损失小于1%

加速比（Speedup Ratio）

推理速度提升的倍数。实际加速比可能与理论计算量减少不成正比，因为还存在内存访问、框架开销等因素。

内存占用减少（Memory Reduction）

推理过程中内存占用的减少量。这个指标对于边缘设备部署尤为重要。

2.2 量化技术

2.2.1 量化原理

量化（Quantization）是将连续取值（或高精度的离散取值）的数据映射到低精度的离散取值空间的过程。在深度学习中，量化特指将模型参数和计算从高精度浮点数（如FP32）转换为低精度整数（如INT8、INT4甚至INT2）的技术。

为什么量化可以压缩模型？

FP32浮点数占用4字节（32位），而INT8整数占用1字节（8位）。从FP32到INT8，理论上可以将模型体积减少4倍。同时，现代CPU/GPU的整数运算单元比浮点运算单元更高效，可以实现3-4倍的推理加速。

量化的数学表示

设原始浮点数据为\(r\)，量化的整数值为\(q\)，量化参数包括缩放因子\(S\)（Scale）和零点\(Z\)（Zero Point）。

非对称量化的转换公式为：

\[q = \text{round}\left(\frac{r}{S}\right) + Z\]

对应的反量化公式为：

\[r = S \times (q - Z)\]

对称量化的零点\(Z=0\)，公式简化为：

\[q = \text{round}\left(\frac{r}{S}\right)\]

\[r = S \times q\]

量化参数的选择

量化参数\(S\)和\(Z\)的选取直接决定了量化后的数值分布和精度损失。常见的参数选择方法包括：

MinMax方法：基于数据的实际范围计算\(S\)和\(Z\)

\[S = \frac{r_{\max} - r_{\min}}{q_{\max} - q_{\min}}\]

\[Z = q_{\max} - \frac{r_{\max}}{S}\]

KL散度方法：通过最小化量化前后的分布差异来选择参数（TensorRT使用这种方法）
阈值校准方法：选择一个阈值\(T\)截断极端值，只对\([-\alpha, \alpha]\)范围内的数据进行量化

2.2.2 对称量化与非对称量化

对称量化（Symmetric Quantization）的零点固定为0，数据分布在0点两侧。在深度学习中，权重（Weight）的分布通常关于0对称，适合使用对称量化。

对称量化的数值范围为\([-R, R]\)，其中\(R\)是绝对值最大的元素。缩放因子计算为：

\[S = \frac{R}{2^{N-1}}\]

其中\(N\)是量化位数。对于INT8，\(N=8\)，所以\(S = R / 128\)。

对称量化的优势是实现简单，硬件处理方便；劣势是不能很好地处理非对称分布的数据。

非对称量化（Asymmetric Quantization）的零点\(Z\)可以不为0，可以更好地适应非对称分布的数据。激活值（Activation）的分布通常不是关于0对称的，适合使用非对称量化。

非对称量化的数值范围为\([r_{\min}, r_{\max}]\)，对应的缩放因子为：

\[S = \frac{r_{\max} - r_{\min}}{2^N - 1}\]

零点为：

\[Z = -\text{round}\left(\frac{r_{\min}}{S}\right)\]

非对称量化的优势是能更充分利用量化后的数值范围，精度损失更小；劣势是实现稍复杂，需要处理零点偏移。

实际选择建议：

权重（Weight）：通常使用对称量化，因为权重分布接近对称
激活值（Activation）：通常使用非对称量化，因为ReLU等激活函数会导致分布偏移
某些特殊场景（如BERT中的LayerNorm）：可能需要混合使用两种量化方式

2.2.3 饱和量化与非饱和量化

在将浮点数映射到整数时，有两种策略：饱和量化（Saturating Quantization）和非饱和量化（Non-saturating Quantization）。

非饱和量化直接按照数值的实际比例进行映射，超出范围的数值会被截断：

\[q = \text{clamp}\left(\text{round}\left(\frac{r}{S}\right), q_{\min}, q_{\max}\right)\]

如果数据分布中存在极端值（Outlier），这些极端值会被截断，可能导致信息丢失。

饱和量化先用一个阈值\(T\)截断超出范围的极端值，再进行量化：

\[r' = \text{clip}(r, -T, T)\]

\[q = \text{round}\left(\frac{r'}{S}\right)\]

通过截断极端值，可以让剩余大部分数据在量化后有更好的表示精度。对于存在噪声或异常值的数据集，饱和量化通常表现更好。

阈值\(T\)的选择是饱和量化的关键。常用的方法是KL散度校准：遍历不同的阈值，计算原始分布和量化分布之间的KL散度，选择使KL散度最小的阈值作为\(T\)。

2.2.4 量化粒度

量化可以在不同的粒度上进行：逐张量量化（Per-Tensor）和逐通道量化（Per-Channel）。

逐张量量化对整个张量使用统一的量化参数。所有元素共享同一个缩放因子\(S\)和零点\(Z\)。这是最简单的量化方式，实现开销最小，但可能因为不同通道/神经元数值分布差异大而导致精度损失。

逐通道量化对张量的不同通道（如卷积的输出通道）使用不同的量化参数。每个通道有独立的\(S\)和\(Z\)。逐通道量化能更好地捕捉不同通道的数值分布差异，通常能获得更好的精度，但实现开销稍大。

对于卷积神经网络：

权重（Weight）：逐通道量化（Per-Channel），每个输出通道有独立的量化参数
激活值（Activation）：逐张量量化（Per-Tensor）或逐token量化（Per-Token）

2.3 训练后量化

2.3.1 训练后量化概述

训练后量化（Post-Training Quantization，PTQ）是在模型训练完成之后，对模型参数进行量化的一种方法。PTQ不需要重新训练模型，只需要少量校准数据来确定量化参数，因此实施成本低、周期短。

PTQ的基本流程：

加载预训练模型：读取FP32精度的训练好的模型
准备校准数据：收集一小部分有代表性的输入数据（通常500-1000个样本）
执行模型推理：用校准数据运行模型前向传播，收集各层输出的数值分布
计算量化参数：根据数值分布确定每层的量化参数（\(S\)和\(Z\)）
执行量化转换：将FP32参数转换为INT8表示
验证量化精度：在测试集上评估量化后模型的精度

PTQ的优势是简单易行，不需要GPU训练资源；劣势是精度损失可能较大，尤其是对于激活值分布复杂或存在极端值的模型。

2.3.2 动态量化与静态量化

动态量化（Dynamic Quantization）的量化参数在推理时动态计算。具体来说：

权重在模型加载时就已经量化好（离线量化）
激活值在每次推理时根据实际输入数据的分布动态计算量化参数

动态量化的优势是精度损失小（因为激活值的量化参数是根据实际输入计算的）；劣势是每次推理都需要计算量化参数，引入额外开销。

静态量化（Static Quantization）的量化参数在模型加载前就预先确定。具体来说：

权重和激活值都在离线阶段完成量化
推理时直接使用预先计算好的量化参数

静态量化的优势是推理速度快（不需要计算量化参数）；劣势是精度损失可能较大（因为激活值的量化参数是基于校准数据而非实际输入）。

对于不同类型的模型：

BERT等Transformer模型：推荐使用动态量化，因为激活值分布复杂且输入变化大
ResNet等CNN模型：可以使用静态量化，因为激活值分布相对稳定

2.3.3 量化参数校准

量化参数校准是PTQ中最关键的步骤，直接决定量化精度。

基于数值范围的校准（MinMax Calibration）

最简单的校准方法，基于数据的实际数值范围：

\[S = \frac{r_{\max} - r_{\min}}{q_{\max} - q_{\min}}\]

\[Z = q_{\max} - \frac{r_{\max}}{S}\]

这种方法简单但容易受到极端值影响。如果数据中存在少量极端值，数值范围会被拉大，导致大部分数据的量化精度下降。

基于直方图的校准（Histogram Calibration）

将数据分布整理为直方图，选择一个最优阈值进行截断和量化：

收集数据的直方图分布
遍历不同的阈值\(T\)，计算截断后的KL散度
选择使KL散度最小的\(T\)作为截断阈值
根据阈值计算量化参数

KL散度校准是TensorRT等推理引擎使用的方法，通常能获得比MinMax更好的精度。

基于阈值搜索的校准

对于某些特殊的激活函数（如ReLU6），可能需要特殊的处理。通过搜索不同的截断范围，选择使量化后模型精度最高或量化误差最小的参数。

2.3.4 量化实践流程

以下是一个典型的PTQ实践流程，以PyTorch为例：

第一步：加载预训练模型

import torch
import torchvision.models as models

# 加载FP32预训练模型
model = models.resnet50(pretrained=True)
model.eval()

第二步：准备校准数据

# 准备校准数据集
calibration_data = []
for images, _ in torch.utils.data.DataLoader(
    torchvision.datasets.CIFAR10('./data', train=False, download=True),
    batch_size=32
):
    calibration_data.append(images)
    if len(calibration_data) >= 16:  # 使用16个batch作为校准数据
        break

第三步：执行动态量化

# 对模型权重执行动态量化
# dtype=quint8表示INT8无符号整数，qconfig指定量化配置
model_quantized = torch.quantization.quantize_dynamic(
    model,  # 原始模型
    {torch.nn.Linear, torch.nn.Conv2d},  # 需要量化的层类型
    dtype=torch.qint8  # 量化数据类型
)

第四步：验证量化精度

# 在测试集上评估量化后模型
correct = 0
total = 0
with torch.no_grad():
    for images, labels in test_loader:
        outputs = model_quantized(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'量化后精度: {100 * correct / total:.2f}%')

2.4 量化感知训练

2.4.1 量化感知训练概述

量化感知训练（Quantization-Aware Training，QAT）是在模型训练过程中模拟量化效应的一种方法。通过在训练阶段引入量化操作，让模型"适应"量化带来的误差，从而在量化后保持较好的精度。

PTQ的局限在于：模型训练时使用的是FP32精度，参数已经"习惯"了高精度表示。当突然转换为低精度时，精度损失可能较大。

QAT的核心思想是"让模型提前知道量化会发生什么"。在训练阶段，模拟量化引入的噪声，让参数在更新时考虑量化效应。这样，当实际部署进行量化时，精度损失会更小。

2.4.2 伪量化节点

QAT通过在计算图中插入伪量化节点（FakeQuant Node）来模拟量化效应。

伪量化节点的工作原理

在正向传播中，伪量化节点执行"量化-反量化"操作：

def fake_quant(x, scale, zero_point, num_bits=8):
    # 量化
    qx = torch.clamp(torch.round(x / scale) + zero_point, 
                     0, 2**num_bits - 1)
    # 反量化
    return (qx - zero_point) * scale

虽然这个操作在数学上应该恒等于输入，但由于浮点舍入误差，实际上引入了量化噪声。

在反向传播中，梯度需要近似传递。由于四舍五入操作的梯度在大多数点为0或1，QAT使用直通估计器（Straight-Through Estimator，STE）来近似梯度：

\[\frac{\partial \text{FakeQuant}(x)}{\partial x} \approx 1\]

STE允许梯度畅通无阻地通过伪量化节点，使参数能够正常更新。

2.4.3 QAT训练流程

QAT的训练流程如下：

准备FP32预训练模型：通常基于已训练好的FP32模型开始
插入伪量化节点：在需要量化的层后插入FakeQuant节点
训练/微调：使用正常训练流程进行训练
移除伪量化节点：在实际部署前，将伪量化节点的属性融合到参数中
导出量化模型：导出量化后的模型用于部署

QAT的训练配置建议

学习率：使用较小的学习率（原始训练学习率的1%）
训练周期：通常只需要原始训练周期的5-10%
优化器：推荐使用带动量的SGD，而非Adam/RMSProp
学习率调度：使用余弦退火（Cosine Annealing）

2.4.4 QAT与PTQ对比

特性	PTQ	QAT
实施难度	简单，无需训练	复杂，需要训练/微调
精度损失	较大（通常2-5%）	较小（通常<1%）
实施周期	短（数小时）	长（数天）
计算资源	少量校准数据	GPU训练资源
适用场景	快速验证、CNN模型	高精度要求、Transformer模型

选择建议：

如果时间紧迫或没有训练资源，选择PTQ
如果对精度要求高，选择QAT
CNN模型通常PTQ就能获得可接受的精度
Transformer模型（如BERT）建议使用QAT

2.5 模型剪枝

2.5.1 剪枝概述

模型剪枝（Pruning）通过删除模型中"不重要"的参数来减小模型规模。与量化通过改变数值精度来压缩模型不同，剪枝直接减少参数数量。

剪枝的理论基础是神经网络的过参数化（Over-parameterization）现象。研究表明，很多训练好的神经网络中存在大量冗余的参数。这些参数对模型输出的贡献很小甚至没有贡献，可以被安全地删除而不显著影响模型精度。

剪枝的形式化描述：给定原始参数向量\(\mathbf{w}\)和目标稀疏度\(k\)，找到一个稀疏参数向量\(\mathbf{w}'\)满足：

\[\|\mathbf{w}'\|_0 \leq k\]

使得在验证集上的性能损失最小。

2.5.2 非结构化剪枝

非结构化剪枝（Unstructured Pruning）以单个参数为单位进行剪枝，即随机地将某些权重设为0。这种方法可以达到很高的压缩率（稀疏度），但需要特殊的硬件支持才能实际加速推理。

基于权重大小的剪枝

最简单的方法是：根据权重的绝对值大小判断其重要性。绝对值越小的权重，对模型输出的影响越小，越容易被剪枝。

剪枝过程：

设定目标稀疏度（如50%）
计算所有权重的绝对值
将绝对值最小的50%权重设为0
保留其余50%的权重不变

剪枝与微调

剪枝后模型精度会下降，需要通过微调恢复。典型的流程是：

训练一个过参数化的大模型
执行剪枝，得到稀疏模型
在训练数据上微调稀疏模型
可能需要多次迭代：剪枝→微调→剪枝→微调...

非结构化剪枝的局限

虽然非结构化剪枝可以达到很高的稀疏度，但实际推理加速有限。原因：

大多数硬件（CPU/GPU）针对密集矩阵运算优化，稀疏表示反而可能更慢
稀疏矩阵需要额外的索引数据结构，增加了存储和计算开销
需要专门的稀疏矩阵运算库（如cuSPARSE）才能获得实际加速

2.5.3 结构化剪枝

结构化剪枝（Structured Pruning）以结构为单位（如整个神经元、卷积核、通道）进行剪枝。删除的是一个完整的结构单元，而非单个参数。

结构化剪枝的优势是：

剪枝后的模型仍然是密集的，可以直接使用标准矩阵运算库加速
不需要特殊的硬件或软件支持
实际加速效果与理论压缩率接近

卷积核剪枝（Kernel/Filter Pruning）

删除整个卷积核是最常见的结构化剪枝形式。设原始卷积层有\(N\)个卷积核，剪枝后保留\(M < N\)个卷积核。

卷积核剪枝的影响：

输出特征图的通道数减少
下一层的计算量和参数量也相应减少
模型精度可能下降，需要微调恢复

通道剪枝（Channel Pruning）

通道剪枝删除特征图的某个通道，等价于删除所有使用该通道作为输入的卷积核。

通道剪枝的影响：

特征图的空间尺寸不变，通道数减少
与该通道相连的所有卷积核权重都需要删除
下一层的输入通道数减少

层级剪枝（Layer Pruning）

直接删除整个网络层（如某个卷积层或某个Transformer层）。这是最激进的剪枝方式，可能带来显著的加速和压缩，但也可能导致精度大幅下降。

2.5.4 剪枝算法

基于重要性的剪枝

最简单的方法是基于某种重要性指标对结构进行排序，然后剪掉最不重要的部分。常用指标包括：

权重的L1范数：\(\|w\|_1 = \sum |w_i|\)，L1范数小的权重被认为不重要
权重的L2范数：\(\|w\|_2 = \sqrt{\sum w_i^2}\)，类似L1
BatchNorm的缩放因子\(\gamma\)：训练时对通道添加缩放因子，\(\gamma\)小的通道被剪枝
Taylor展开近似：近似计算移除参数对损失函数的影响

基于正则化的剪枝

在训练损失函数中添加稀疏正则项（如L1范数），促使模型自动学习到稀疏结构。训练完成后，删除接近零的参数。

基于重构误差的剪枝

最小化剪枝前后的特征重构误差。例如，对于卷积层的通道剪枝，选择删除后特征重构误差最小的通道组合。

** Lottery Ticket Hypothesis**

MIT的研究者提出了"Lottery Ticket Hypothesis"：一个随机初始化的密集网络包含一个子网络（"中奖彩票"），该子网络单独训练可以达到与原始网络相同的精度，且收敛速度更快。

剪枝过程就是寻找这个"中奖彩票"的过程：

随机初始化一个密集网络
训练到收敛
剪枝得到子网络（保留原始初始化权重）
单独训练子网络

如果子网络能达到原始网络相近的精度，说明找到了"中奖彩票"。

2.6 知识蒸馏

2.6.1 知识蒸馏概述

知识蒸馏（Knowledge Distillation，KD）是由Hinton等人在2015年提出的概念，核心思想是让"学生模型"学习"教师模型"的知识来实现模型压缩。

传统的模型压缩方法（如剪枝、量化）主要关注模型参数本身，而知识蒸馏关注的是模型知识的传递。知识可以被理解为模型对数据的表示能力和泛化能力。

知识蒸馏的框架：

教师模型（Teacher Model）：大而精确的模型，通常是训练好的原始模型
学生模型（Student Model）：小而高效的模型，是蒸馏后希望得到的压缩模型
蒸馏温度（Temperature）：控制软标签分布的"软化"程度

2.6.2 蒸馏知识类型

教师模型可以传递给学生模型的知识有多种类型：

基于响应的知识（Response-Based Knowledge）

最直接的知识形式是教师模型的输出（软标签）。对于分类问题，教师模型输出的各类别概率分布包含了比硬标签更丰富的信息。

例如，对于一张猫狗分类的图片：

硬标签：[1, 0]（假设"狗"=1，"猫"=0）
软标签：[0.7, 0.3]（教师模型预测的概率）

软标签中的0.3概率说明模型认为这张图片有30%的可能性是猫，即使最终分类为狗。这种"不确定"信息反映了模型的认知，对学生学习更有价值。

基于特征的知识（Feature-Based Knowledge）

教师模型的中间层输出也包含知识。学生模型可以学习教师模型中间层的特征表示。

对于CNN，中间层特征通常是特征图（Feature Map）；对于Transformer，是注意力权重或隐藏状态。

学生模型不一定与教师模型有相同的结构，可以通过适配层（Adaptor）将教师特征转换为学生能学习的目标形式。

基于关系知识（Relation-Based Knowledge）

超越单个输出或特征，学习不同输出/特征之间的关系。例如：

教师模型不同层输出之间的关系（如Gram矩阵）
数据样本在教师模型特征空间中的相似度关系
教师模型对不同样本的预测顺序关系

2.6.3 蒸馏算法

经典知识蒸馏（Classic Distillation）

Hinton等人提出的经典算法，流程如下：

用训练数据同时训练教师模型和学生模型
教师模型生成软标签（通过高温度\(T\)的softmax）
学生模型同时学习硬标签和软标签

损失函数为：

\[\mathcal{L} = \alpha \mathcal{L}_{soft} + (1-\alpha) \mathcal{L}_{hard}\]

其中：

\(\mathcal{L}_{soft}\)：学生与教师软标签的交叉熵
\(\mathcal{L}_{hard}\)：学生与真实标签的交叉熵
\(\alpha\)：两类损失的权重

温度参数\(T\)的作用

标准softmax将logits转换为概率：

\[q_i = \frac{e^{z_i/T}}{\sum_j e^{z_j/T}}\]

当\(T=1\)时，就是标准softmax。当\(T > 1\)时，概率分布变得更"软"（Smooth），各类别之间的概率差异被缩小。

高温下的软标签能更好地揭示类别之间的相似性。例如，"狗"和"狼"在高温下概率会比较接近，而"狗"和"汽车"则相差很大。

在线蒸馏（Online Distillation）

在线蒸馏中，教师和学生模型同时训练、同时更新。教师模型可以是多个模型的ensemble，也可以是和学生模型相同结构的多个模型。

优势是不需要预训练教师模型；劣势是教师模型的质量可能不如预先训练的好。

自蒸馏（Self-Distillation）

自蒸馏是在线蒸馏的特殊形式：教师模型和学生模型结构相同，知识从深层/大模型传递到浅层/小模型。

典型应用：

深层特征指导浅层特征学习
网络整体指导网络的一部分（如ResNet的 distillation from deeper layers to shallower layers）
训练过程中的自己指导自己（如Label Smoothing的变体）

2.6.4 蒸馏实践

蒸馏训练流程

准备教师模型：通常使用在大数据集上预训练好的高精度模型
准备学生模型：设计比教师模型更简单的结构
准备训练数据：可以使用与训练教师相同的数据集
设置蒸馏温度：通常\(T=2-20\)，需要调试选择
执行蒸馏训练：同时用硬标签和软标签监督学生模型
评估学生模型：在测试集上评估蒸馏后模型的精度

蒸馏训练示例

import torch
import torch.nn as nn
import torch.nn.functional as F

class DistillationLoss(nn.Module):
    def __init__(self, temperature=4.0, alpha=0.5):
        super().__init__()
        self.temperature = temperature
        self.alpha = alpha

    def forward(self, student_logits, teacher_logits, labels):
        # 硬标签损失
        hard_loss = F.cross_entropy(student_logits, labels)

        # 软标签损失（知识蒸馏）
        soft_teacher = F.softmax(teacher_logits / self.temperature, dim=1)
        soft_student = F.log_softmax(student_logits / self.temperature, dim=1)
        soft_loss = F.kl_div(soft_student, soft_teacher, reduction='batchmean')
        soft_loss = soft_loss * (self.temperature ** 2)  # 补偿温度因子的缩放

        # 组合损失
        return self.alpha * soft_loss + (1 - self.alpha) * hard_loss

# 蒸馏训练循环
for epoch in range(num_epochs):
    for batch in train_loader:
        inputs, labels = batch

        # 教师模型前向（不更新梯度）
        with torch.no_grad():
            teacher_logits = teacher_model(inputs)

        # 学生模型前向
        student_logits = student_model(inputs)

        # 计算蒸馏损失
        loss = distillation_loss(student_logits, teacher_logits, labels)

        # 更新学生模型
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

蒸馏温度的选择

温度\(T\)的选择影响软标签的分布：

\(T=1\)：标准softmax，不做任何"软化"
\(T=2-5\)：轻度软化，类别间差异仍然明显
\(T=10-20\)：高度软化，几乎所有类别都有一定概率

建议从\(T=4\)开始尝试，根据学生模型精度调整。

2.7 模型压缩综合应用

2.7.1 量化与剪枝的结合

量化和剪枝是从不同维度压缩模型的方法，可以组合使用以获得更高的压缩率。

先剪枝后量化

一种常见的流程是：先通过剪枝得到稀疏模型，再对稀疏模型进行量化。由于剪枝已经删除了大量参数，量化时需要处理的数据量更少。

优势：

可以分别发挥两种方法的优势
剪枝后模型更小，量化参数更少，校准更简单

注意事项：

剪枝可能改变参数分布，影响量化效果
某些剪枝模式（如结构化剪枝）可能更适合后续量化

联合优化

更先进的方法是同时优化剪枝和量化。例如，在训练损失中加入稀疏性正则项和量化误差项，让模型同时学习稀疏结构和量化表示。

2.7.2 知识蒸馏与量化的结合

知识蒸馏和量化也可以结合使用：

量化感知蒸馏（Quantization-Aware Distillation）

在蒸馏训练过程中引入量化感知，让学生模型学习如何在量化约束下保持精度。具体做法是在蒸馏训练时给教师模型也插入伪量化节点，让教师模拟量化后的行为。

蒸馏补偿量化误差

蒸馏可以用于补偿量化引入的精度损失。思路是：先用量化得到一个低精度模型，再用蒸馏让高精度教师指导这个低精度模型学习，恢复精度。

本章小结

模型压缩目的：在保持模型性能的前提下减小存储体积、降低计算复杂度和减少内存占用，主要方法包括量化、剪枝、知识蒸馏和低秩分解。
量化技术：通过将FP32参数转换为INT8等低精度表示来压缩模型。核心概念包括对称/非对称量化、饱和/非饱和量化、逐张量/逐通道量化等。PTQ实施简单但精度损失可能较大，QAT通过训练模拟量化效应能获得更好的精度。
模型剪枝：通过删除"不重要"的参数来减小模型规模。非结构化剪枝粒度细但硬件支持有限，结构化剪枝粒度粗但可直接加速。常用方法包括基于重要性、基于正则化、基于重构误差等。
知识蒸馏：通过让学生模型学习教师模型的知识来实现压缩。蒸馏知识可以是基于响应的、基于特征的或基于关系的。经典蒸馏使用软标签和硬标签联合训练，温度参数控制软化程度。
综合应用：量化和剪枝可以组合使用，知识蒸馏也可与量化结合，通过蒸馏补偿量化带来的精度损失。

思考与练习

理解概念：解释PTQ和QAT的核心区别。为什么QAT通常能获得比PTQ更好的量化精度？
原理分析：假设要对一个卷积层的权重进行INT8量化，权重分布直方图如下所示（大部分值集中在0附近，少量值较大）。请分析：
应该选择对称量化还是非对称量化？为什么？
如果使用饱和量化，截断阈值应该如何选择？
设计思考：如果要设计一个针对MobileNet等轻量级网络的剪枝策略，请说明：
会选择结构化剪枝还是非结构化剪枝？为什么？
如何确定每个层的剪枝率？
实践应用：在知识蒸馏中，温度参数\(T\)的作用是什么？\(T\)过大或过小会有什么影响？
综合分析：假设你需要在边缘设备上部署一个目标检测模型，该设备只有INT8硬件加速支持，且内存有限。请设计一个综合使用量化、剪枝和知识蒸馏的压缩方案。
扩展思考：近年来，大语言模型（LLM）的模型压缩成为研究热点。请调研并分析LLM quantization（如GPTQ、AWQ）与传统量化方法的主要区别和挑战。