知识蒸馏与模型压缩深度解析（Knowledge Distillation & Model Compression）
·
目录
1. 模型压缩概述与理论基础
1.1 为什么需要模型压缩
模型压缩的动机:
现代深度学习模型越来越大的趋势:
┌─────────────────────────────────────────────────────────────────┐
│ 模型规模增长 │
├───────────────┬───────────┬─────────────────────────────────────┤
│ 模型 │ 参数量 │ 推理需求 │
├───────────────┼───────────┼─────────────────────────────────────┤
│ ResNet-50 │ 25M │ 4 GFLOPs, ~10ms (GPU) │
│ BERT-Large │ 340M │ 60 GFLOPs, ~50ms │
│ GPT-3 │ 175B │ 314 TFLOPs, ~秒级 │
│ GPT-4 │ ~1.8T │ ~3 TFLOPs/token │
│ LLaMA-70B │ 70B │ 140 GFLOPs/token │
└───────────────┴───────────┴─────────────────────────────────────┘
问题:
1. 推理延迟: 大模型推理慢，无法满足实时需求
2. 内存占用: 大模型需要大量显存/内存
3. 部署成本: 需要昂贵的 GPU 集群
4. 能耗: 大模型推理功耗高
5. 边缘部署: 移动设备、嵌入式设备资源有限
模型压缩的目标:
在尽量保持性能的前提下，减小模型大小和计算量
1.2 压缩方法分类
┌─────────────────────────────────────────────────────────────────────┐
│ 模型压缩方法分类 │
├─────────────────────────────────────────────────────────────────────┤
│ │
│ 1. 知识蒸馏 (Knowledge Distillation) │
│ └─ 训练小模型模仿大模型的行为 │
│ │
│ 2. 模型剪枝 (Pruning) │
│ └─ 移除不重要的权重/神经元/层 │
│ ├── 非结构化剪枝: 移除单个权重 │
│ └── 结构化剪枝: 移除整个通道/头/层 │
│ │
│ 3. 模型量化 (Quantization) │
│ └─ 降低权重和激活的数值精度 │
│ ├── 训练后量化 (PTQ) │
│ └── 量化感知训练 (QAT) │
│ │
│ 4. 低秩分解 (Low-Rank Factorization) │
│ └─ 将大矩阵分解为小矩阵的乘积 │
│ │
│ 5. 神经架构搜索 (NAS) │
│ └─ 自动设计高效的网络架构 │
│ │
│ 6. 稀疏化 (Sparsification) │
│ └─ 使模型权重稀疏，减少计算 │
│ │
└─────────────────────────────────────────────────────────────────────┘
1.3 压缩效果评估
评估指标:
1. 模型大小 (Model Size):
参数量 × 每个参数的字节数
FP32: 4 bytes/param
FP16: 2 bytes/param
INT8: 1 byte/param
INT4: 0.5 bytes/param
2. 计算量 (FLOPs):
浮点运算次数
通常用 MAC (乘加) 衡量: 1 MAC ≈ 2 FLOPs
3. 推理延迟 (Latency):
单次推理时间
取决于硬件、实现、batch size
4. 吞吐量 (Throughput):
每秒处理的样本数
5. 内存占用 (Memory):
模型权重 + 激活值 + KV Cache
6. 准确率保持:
压缩后模型的性能下降程度
压缩比:
压缩比 = 原始模型大小 / 压缩后模型大小
例:
ResNet-50 (FP32, 100MB) → INT8 (25MB): 4x 压缩
LLaMA-7B (FP16, 14GB) → INT4 (3.5GB): 4x 压缩
2. 知识蒸馏核心理论
2.1 知识蒸馏框架
知识蒸馏 (Knowledge Distillation, KD) 的核心思想:
训练一个小模型 (Student) 来模仿一个大模型 (Teacher) 的行为
┌─────────────────────────────────────────────────────────────────┐
│ │
│ Teacher Model (大) Student Model (小) │
│ ┌──────────────┐ ┌──────────────┐ │
│ │ │ │ │ │
│ │ 340M params │ 知识传递 │ 66M params │ │
│ │ │ ──────────────► │ │ │
│ │ 高准确率 │ │ 接近的准确率 │ │
│ │ │ │ │ │
│ └──────────────┘ └──────────────┘ │
│ │
│ 理论: Teacher 的 "软标签" 包含比硬标签更多的信息 │
│ │
└─────────────────────────────────────────────────────────────────┘
2.2 Hinton 知识蒸馏
论文: "Distilling the Knowledge in a Neural Network" (Hinton et al., 2015)
核心贡献:
1. 软标签 (Soft Labels): 使用 Teacher 的 softmax 输出作为监督信号
2. 温度缩放 (Temperature): 调整 softmax 的温度，揭示更多信息
3. 混合损失: 结合硬标签和软标签
理论基础:
Teacher 的 softmax 输出:
p_i = exp(z_i/T) / Σ_j exp(z_j/T)
其中:
z_i: Teacher 的 logits
T: 温度参数
温度的作用:
T=1: 标准 softmax，分布尖锐
T>1: 分布更平滑，揭示类间关系
T→∞: 均匀分布
T→0: one-hot (硬标签)
理论: 高温 softmax 暴露了 Teacher 学到的"暗知识"
(dark knowledge) — 类别之间的相似性
例: 猫的图像
T=1: [0.9, 0.05, 0.05] (猫, 狗, 鸟)
T=5: [0.5, 0.3, 0.2]
高温揭示: 猫和狗比猫和鸟更相似
损失函数:
L = α · T² · KL(p_Teacher ‖ p_Student) + (1-α) · CE(y_hard, p_Student)
其中:
KL: KL 散度 (软标签损失)
CE: 交叉熵 (硬标签损失)
α: 平衡系数 (通常 0.5-0.9)
T²: 温度平方缩放 (补偿梯度缩放)
import torch
import torch.nn as nn
import torch.nn.functional as F
class KnowledgeDistillationLoss(nn.Module):
"""
Hinton 知识蒸馏损失
L = α · T² · KL(p_T ‖ p_S) + (1-α) · CE(y, p_S)
理论:
软标签: Teacher 的高温 softmax 输出
硬标签: 真实标签
软标签包含更多信息:
- 类别间相似性
- 样本难度
- 不确定性
"""
def __init__(self, temperature=4.0, alpha=0.7):
super().__init__()
self.T = temperature
self.alpha = alpha
def forward(self, student_logits, teacher_logits, hard_labels):
"""
student_logits: Student 模型的 logits [B, C]
teacher_logits: Teacher 模型的 logits [B, C]
hard_labels: 真实标签 [B]
"""
# 软标签损失 (KL 散度)
p_teacher = F.softmax(teacher_logits / self.T, dim=-1)
p_student = F.log_softmax(student_logits / self.T, dim=-1)
# KL 散度
soft_loss = F.kl_div(
p_student,
p_teacher,
reduction='batchmean'
) * (self.T ** 2)
# 硬标签损失 (交叉熵)
hard_loss = F.cross_entropy(student_logits, hard_labels)
# 混合损失
total_loss = self.alpha * soft_loss + (1 - self.alpha) * hard_loss
return total_loss
"""
温度 T 的影响:
T=1: 标准 softmax，类间差异大
Student 学到的信息有限
T=2-4: 适中的温度，推荐值
揭示类间关系，但不过于平滑
T=10+: 高温，分布过于平滑
可能丢失有用信息
理论最优温度:
取决于 Teacher 的置信度
置信度高 → 需要更高温度
置信度低 → 较低温度即可
"""
2.3 软标签的理论分析
为什么软标签比硬标签更好?
信息论视角:
硬标签: one-hot 向量，熵 = 0 (确定性)
软标签: 概率分布，熵 > 0 (包含不确定性信息)
软标签的信息量更大
梯度分析:
硬标签的梯度:
∂CE/∂z_i = p_i - y_i
对于正确类别: 梯度 = p_correct - 1
对于错误类别: 梯度 = p_wrong
软标签的梯度:
∂KL/∂z_i = (p_i^S - p_i^T) / T
梯度与 Student 和 Teacher 的差异成正比
含义:
硬标签: 只告诉 Student "对不对"
软标签: 告诉 Student "差多少"
类间关系:
硬标签: 猫 ≠ 狗 ≠ 鸟 (完全独立)
软标签: 猫 ≈ 狗 ≠ 鸟 (编码相似性)
这种"暗知识"是 Teacher 从数据中学到的
温度缩放的梯度效应:
∂L/∂z_i = (1/T) · (p_i^S - p_i^T)
当 T 大时:
p_i^T 更平滑
梯度更均匀
Student 学习更稳定
2.4 Teacher 的选择
Teacher 模型的选择:
1. 单个强 Teacher:
使用性能最好的大模型
优点: 简单，效果好
缺点: 单一视角
2. 集成 Teacher (Ensemble):
使用多个 Teacher 的平均输出
p_T = (1/K) Σ_k p_T_k
理论:
集成减少了单个 Teacher 的偏差
提供更准确的软标签
Student 可以学到多个 Teacher 的优点
3. 助教 Teacher (Assistant):
Teacher → Assistant → Student
理论:
当 Teacher 和 Student 差距太大时
直接蒸馏效果不好
引入中等大小的 Assistant 作为桥梁
4. 自蒸馏 (Self-Distillation):
模型蒸馏自己 (或自己的 EMA)
理论:
不需要额外的 Teacher
通过数据增强和 EMA 产生不同的"视图"
3. 知识蒸馏变体与进阶
3.1 特征蒸馏（Feature Distillation）
理论动机:
不仅匹配输出，还匹配中间特征
软标签蒸馏: 只匹配最终输出
特征蒸馏: 匹配中间层的特征表示
方法:
FitNets (Romero et al., 2015):
匹配 Teacher 和 Student 的中间层特征
L_feat = Σ_l MSE(f_T^l, transform(f_S^l))
其中 transform 是适配层 (当维度不匹配时)
数学推导:
设 Teacher 的第 l 层特征: f_T^l ∈ ℝ^{d_T}
设 Student 的第 l 层特征: f_S^l ∈ ℝ^{d_S}
当 d_T ≠ d_S 时，需要适配层:
f̂_S^l = W_l · f_S^l + b_l
损失:
L_feat = ‖f_T^l - f̂_S^l‖²
class FeatureDistillationLoss(nn.Module):
"""
特征蒸馏损失
匹配 Teacher 和 Student 的中间层特征
理论:
中间层特征包含更丰富的表示信息
比仅匹配输出更有指导性
"""
def __init__(self, teacher_dims, student_dims):
super().__init__()
# 适配层 (当维度不匹配时)
self.adaptors = nn.ModuleList([
nn.Linear(s_dim, t_dim)
for s_dim, t_dim in zip(student_dims, teacher_dims)
])
def forward(self, teacher_features, student_features):
"""
teacher_features: List[Tensor], Teacher 各层特征
student_features: List[Tensor], Student 各层特征
"""
loss = 0
for i, (t_feat, s_feat) in enumerate(zip(teacher_features, student_features)):
# 适配 Student 特征
s_feat_adapted = self.adaptors[i](s_feat)
# MSE 损失
loss += F.mse_loss(s_feat_adapted, t_feat)
return loss / len(teacher_features)
"""
特征蒸馏的层选择:
理论:
不同层的特征包含不同级别的信息
浅层: 边缘、纹理等低级特征
中层: 部件、模式等中级特征
深层: 语义、类别等高级特征
策略:
1. 全层匹配: 匹配所有层 (计算量大)
2. 关键层匹配: 只匹配几个关键层
3. 渐进式: 从浅到深逐步匹配
"""
3.2 注意力蒸馏（Attention Transfer）
论文: "Paying More Attention to Attention: Improving the Performance of
Convolutional Neural Networks via Attention Transfer" (Zagoruyko & Komodakis, 2017)
核心思想:
匹配 Teacher 和 Student 的注意力图
理论:
注意力图反映模型"关注"的区域
Student 应该关注与 Teacher 相同的区域
注意力图定义:
对于特征图 A ∈ ℝ^{C×H×W}:
通道注意力: A_channel = Σ_c |A_c| (逐通道求和)
空间注意力: A_spatial = Σ_s |A_s| (逐空间位置求和)
或使用更复杂的注意力:
A = g(A) (如全局平均池化)
class AttentionTransferLoss(nn.Module):
"""
注意力蒸馏损失
匹配 Teacher 和 Student 的注意力图
理论:
注意力图 = 特征图的统计量 (如 L2 范数)
Student 应该关注与 Teacher 相同的区域
"""
def __init__(self):
super().__init__()
def compute_attention_map(self, feature_map):
"""
计算注意力图
A = normalize(Σ_c |f_c|²)
"""
# L2 范数作为注意力
attention = torch.sum(feature_map ** 2, dim=1, keepdim=True)
# 归一化
attention = attention / (torch.sum(attention, dim=[2, 3], keepdim=True) + 1e-6)
return attention
def forward(self, teacher_features, student_features):
loss = 0
for t_feat, s_feat in zip(teacher_features, student_features):
t_attn = self.compute_attention_map(t_feat)
s_attn = self.compute_attention_map(s_feat)
loss += F.mse_loss(s_attn, t_attn)
return loss
3.3 关系蒸馏（Relational Knowledge Distillation）
论文: "Relational Knowledge Distillation" (Park et al., 2018)
核心思想:
不仅匹配单个样本的表示，还匹配样本之间的关系
理论:
标准 KD: 匹配 p(y|x) — 单样本
关系 KD: 匹配样本间的相似性结构
对于样本集 {x_i}:
样本间关系: R_ij = similarity(f(x_i), f(x_j))
损失: L = Σ_{i,j} (R_ij^T - R_ij^S)²
关系矩阵:
R_ij = exp(-‖f(x_i) - f(x_j)‖² / σ²)
或使用余弦相似度:
R_ij = cosine(f(x_i), f(x_j))
class RelationalDistillationLoss(nn.Module):
"""
关系蒸馏损失
匹配样本间的关系结构
理论:
不仅匹配绝对值，还匹配相对关系
保持嵌入空间的几何结构
"""
def __init__(self, temperature=1.0):
super().__init__()
self.T = temperature
def compute_relation_matrix(self, features):
"""
计算关系矩阵
R_ij = exp(-‖f_i - f_j‖² / T)
"""
# 计算距离矩阵
dist = torch.cdist(features, features, p=2)
# 转换为相似度
relation = torch.exp(-dist ** 2 / self.T)
return relation
def forward(self, teacher_features, student_features):
# 计算关系矩阵
R_teacher = self.compute_relation_matrix(teacher_features)
R_student = self.compute_relation_matrix(student_features)
# 匹配关系
loss = F.mse_loss(R_student, R_teacher)
return loss
"""
关系蒸馏的优势:
1. 结构保持:
不仅学习单个样本的表示
还学习样本间的几何关系
2. 更强的约束:
N 个样本有 N² 个关系约束
比 N 个独立约束更强
3. 更好的泛化:
关系结构更稳定
对分布偏移更鲁棒
"""
3.4 自蒸馏（Self-Distillation）
理论:
不需要额外的 Teacher
模型从自己或自己的 EMA 中学习
方法:
1. 数据增强自蒸馏:
同一图像的两个增强视图应该有相似的输出
L = KL(p(x_aug1) ‖ p(x_aug2))
2. EMA 自蒸馏 (BYOL 风格):
在线网络: θ (更新)
目标网络: θ' = α·θ' + (1-α)·θ (EMA)
L = MSE(f_θ(x), stop_grad(f_θ'(x)))
3. 深度自蒸馏:
浅层的输出应该接近深层的输出
L = KL(p_shallow ‖ p_deep)
class SelfDistillationLoss(nn.Module):
"""
自蒸馏损失
理论:
使用数据增强和 EMA 产生不同的"视图"
让模型从自己的一致性中学习
"""
def __init__(self, ema_momentum=0.999, temperature=0.1):
super().__init__()
self.momentum = ema_momentum
self.T = temperature
# 目标网络 (EMA)
self.target_model = None
def update_target(self, online_model):
"""更新 EMA 目标网络"""
if self.target_model is None:
self.target_model = copy.deepcopy(online_model)
self.target_model.requires_grad_(False)
else:
# EMA 更新
for p_target, p_online in zip(
self.target_model.parameters(),
online_model.parameters()
):
p_target.data = (self.momentum * p_target.data +
(1 - self.momentum) * p_online.data)
def forward(self, online_features, target_features):
"""
online_features: 在线网络的特征
target_features: 目标网络的特征 (停止梯度)
"""
# L2 归一化
online_norm = F.normalize(online_features, dim=-1)
target_norm = F.normalize(target_features.detach(), dim=-1)
# 余弦相似度
sim = (online_norm * target_norm).sum(dim=-1)
# 损失 (最大化相似度)
loss = -sim.mean()
return loss
3.5 Logit 级蒸馏 vs 特征级蒸馏
┌─────────────────────────────────────────────────────────────────────┐
│ 蒸馏方法对比 │
├───────────────┬──────────────────┬──────────────────────────────────┤
│ 方法 │ 匹配目标 │ 特点 │
├───────────────┼──────────────────┼──────────────────────────────────┤
│ Logit KD │ 输出概率分布 │ 简单，只匹配最终输出 │
│ │ │ 信息量有限 │
├───────────────┼──────────────────┼──────────────────────────────────┤
│ Feature KD │ 中间层特征 │ 信息更丰富 │
│ │ │ 需要适配层 │
├───────────────┼──────────────────┼──────────────────────────────────┤
│ Attention KD │ 注意力图 │ 关注空间信息 │
│ │ │ 适合 CNN │
├───────────────┼──────────────────┼──────────────────────────────────┤
│ Relation KD │ 样本间关系 │ 保持几何结构 │
│ │ │ 更强的约束 │
├───────────────┼──────────────────┼──────────────────────────────────┤
│ Self KD │ 自身一致性 │ 不需要额外 Teacher │
│ │ │ 通过增强和 EMA 实现 │
└───────────────┴──────────────────┴──────────────────────────────────┘
理论分析:
信息量: 特征级 > 注意力级 > Logit 级
实现难度: 特征级 > 注意力级 > Logit 级
通用性: Logit 级 > 特征级 > 注意力级
实践建议:
- 优先使用 Logit KD (简单有效)
- 如果效果不够，加入特征 KD
- 注意力 KD 适合 CNN 任务
4. 模型剪枝
4.1 剪枝概述
模型剪枝 (Pruning) 的核心思想:
移除模型中不重要的部分
减少参数量和计算量
┌─────────────────────────────────────────────────────────────────┐
│ 剪枝类型 │
├─────────────────────────────────────────────────────────────────┤
│ │
│ 非结构化剪枝 (Unstructured Pruning): │
│ 移除单个权重 (设为 0) │
│ 产生稀疏矩阵 │
│ 需要专门的稀疏计算库支持 │
│ 压缩比高，但加速比低 │
│ │
│ 结构化剪枝 (Structured Pruning): │
│ 移除整个通道/卷积核/注意力头/层 │
│ 产生更小但密集的模型 │
│ 不需要专门硬件支持 │
│ 压缩比低，但加速比高 │
│ │
│ 半结构化剪枝 (Semi-Structured Pruning): │
│ N:M 稀疏性 (如 2:4) │
│ 每 M 个权重中有 N 个为零 │
│ 可以利用硬件加速 │
│ │
└─────────────────────────────────────────────────────────────────┘
4.2 重要性评估准则
如何评估权重/结构的重要性?
1. 基于幅度 (Magnitude-based):
重要性 = |w|
理论: 小权重对输出影响小
优点: 简单高效
缺点: 忽略了权重的协同作用
2. 基于梯度 (Gradient-based):
重要性 = |w · ∂L/∂w|
理论: 考虑权重对损失的影响
优点: 更准确
缺点: 需要计算梯度
3. 基于 Hessian (Second-order):
重要性 = w² · H_ii / 2
理论: 二阶近似，更准确
优点: 理论最优
缺点: 计算代价高
4. 基于激活 (Activation-based):
重要性 = E[|activation|]
理论: 不活跃的通道不重要
优点: 直观
缺点: 需要前向传播
class ImportanceScorer:
"""重要性评估器"""
@staticmethod
def magnitude_score(weight):
"""
基于幅度的重要性
重要性 = |w|
理论:
小权重对输出的贡献小
可以安全移除
"""
return weight.abs()
@staticmethod
def gradient_score(weight, gradient):
"""
基于梯度的重要性
重要性 = |w · ∂L/∂w|
理论:
综合考虑权重大小和梯度方向
w 大但梯度小 → 不重要
w 小但梯度大 → 可能重要
"""
return (weight * gradient).abs()
@staticmethod
def taylor_score(weight, gradient):
"""
Taylor 展开的重要性
重要性 = |w · ∂L/∂w + ½ · w² · ∂²L/∂w²|
一阶近似: |w · ∂L/∂w|
"""
return (weight * gradient).abs()
@staticmethod
def activation_score(activations):
"""
基于激活的重要性
重要性 = E[|activation|]
理论:
不活跃的通道/神经元不重要
可以移除
"""
return activations.abs().mean(dim=0)
"""
重要性准则的理论对比:
设真实重要性为 I*(w)，估计重要性为 Î(w)
幅度: Î = |w|
- 误差: 忽略了 w 与其他权重的交互
- 适用: 权重分布近似独立时
梯度: Î = |w · g|
- 误差: 一阶近似
- 适用: 一般情况
Hessian: Î = w · g + ½ · w² · H
- 误差: 二阶近似
- 适用: 需要精确评估时
实践中，幅度准则通常足够好
"""
4.3 剪枝策略
剪枝策略:
1. 一次性剪枝 (One-shot Pruning):
训练完成 → 剪枝 → 微调
优点: 简单
缺点: 高压缩比时性能下降大
2. 迭代剪枝 (Iterative Pruning):
训练 → 剪枝少量 → 微调 → 剪枝 → 微调 → ...
理论:
每次只剪枝少量，模型有时间适应
比一次性剪枝效果好
3. 课程剪枝 (Curriculum Pruning):
从易到难逐步剪枝
理论:
先移除明显不重要的
再逐步移除边缘的
4. 动态剪枝 (Dynamic Pruning):
训练过程中动态调整稀疏度
理论:
不同训练阶段需要不同的稀疏度
早期: 低稀疏度 (探索)
后期: 高稀疏度 (压缩)
class IterativePruner:
"""
迭代剪枝器
策略:
1. 训练模型
2. 评估重要性
3. 剪枝少量 (如 10%)
4. 微调恢复性能
5. 重复步骤 2-4
理论:
每次剪枝少量，模型有时间适应
比一次性剪枝效果更好
"""
def __init__(self, model, target_sparsity=0.9, prune_per_round=0.1):
self.model = model
self.target_sparsity = target_sparsity
self.prune_per_round = prune_per_round
self.current_sparsity = 0
def prune_round(self, train_loader, val_loader):
"""
执行一轮剪枝 + 微调
"""
# 1. 评估重要性
scores = self._compute_importance(train_loader)
# 2. 计算剪枝阈值
current_params = self._count_nonzero()
target_params = int(current_params * (1 - self.prune_per_round))
threshold = self._find_threshold(scores, target_params)
# 3. 应用剪枝
self._apply_pruning(scores, threshold)
# 4. 微调
self._finetune(train_loader, val_loader, epochs=5)
self.current_sparsity += self.prune_per_round
return self.current_sparsity
def _compute_importance(self, data_loader):
"""计算各权重的重要性"""
scores = {}
for name, param in self.model.named_parameters():
if 'weight' in name:
scores[name] = ImportanceScorer.magnitude_score(param.data)
return scores
def _apply_pruning(self, scores, threshold):
"""应用剪枝"""
for name, param in self.model.named_parameters():
if name in scores:
mask = scores[name] > threshold
param.data *= mask.float()
"""
迭代剪枝 vs 一次性剪枝:
一次性剪枝 (压缩 90%):
性能下降: ~5-10%
迭代剪枝 (每轮 10%，共 9 轮):
性能下降: ~1-3%
理论:
每次剪枝后微调，让剩余权重补偿被移除的权重
累积效果更好
"""
4.4 结构化剪枝
class ChannelPruner:
"""
通道剪枝器
移除整个卷积通道
理论:
通道是 CNN 的基本计算单元
移除不重要的通道可以直接减少计算量
不需要特殊硬件支持
"""
def __init__(self, model, prune_ratio=0.3):
self.model = model
self.prune_ratio = prune_ratio
def compute_channel_importance(self, layer):
"""
计算通道重要性
方法 1: L1 范数
重要性 = Σ |w| (对通道的所有权重求和)
方法 2: L2 范数
重要性 = √(Σ w²)
方法 3: BatchNorm 缩放因子
重要性 = |γ| (BN 的缩放参数)
"""
if hasattr(layer, 'weight'):
# 卷积层: 对输出通道求 L1 范数
importance = layer.weight.data.abs().sum(dim=[1, 2, 3])
else:
importance = torch.ones(layer.out_channels)
return importance
def prune_layer(self, layer, keep_ratio):
"""
剪枝单层
1. 计算通道重要性
2. 选择最重要的通道
3. 创建新层
"""
importance = self.compute_channel_importance(layer)
# 选择要保留的通道
n_keep = int(len(importance) * keep_ratio)
_, indices = torch.topk(importance, n_keep)
indices = indices.sort().values
# 创建新层
new_layer = self._create_pruned_layer(layer, indices)
return new_layer
"""
结构化剪枝的挑战:
1. 层间依赖:
移除某层的通道会影响下一层的输入
需要同步剪枝
2. 精度下降:
结构化剪枝的压缩比有限
通常只能剪枝 30-50%
3. 微调需求:
剪枝后需要微调恢复性能
增加训练成本
"""
5. 模型量化
5.1 量化基础理论
量化 (Quantization) 的定义:
将连续的浮点数映射到离散的整数
浮点数: x ∈ ℝ
量化值: x_q ∈ {0, 1, ..., 2^b - 1} (b 位整数)
量化公式:
对称量化:
scale = max(|x|) / (2^(b-1) - 1)
x_q = round(x / scale)
x̂ = x_q × scale
非对称量化:
scale = (max(x) - min(x)) / (2^b - 1)
zero_point = round(-min(x) / scale)
x_q = round(x / scale) + zero_point
x̂ = (x_q - zero_point) × scale
量化误差:
量化引入的误差:
ε = x - x̂ = x - x_q × scale
误差的上界:
|ε| ≤ scale / 2
scale 越小，误差越小 (但需要更多位数)
class SymmetricQuantizer:
"""
对称量化器
理论:
scale = max(|x|) / (2^(b-1) - 1)
x_q = round(x / scale)
x̂ = x_q × scale
特点:
zero_point = 0
对称分布
计算简单
"""
def __init__(self, bits=8):
self.bits = bits
self.qmin = -(2 ** (bits - 1))
self.qmax = 2 ** (bits - 1) - 1
def quantize(self, x):
"""量化"""
# 计算缩放因子
scale = x.abs().max() / self.qmax
# 量化
x_q = torch.clamp(torch.round(x / scale), self.qmin, self.qmax)
# 反量化
x_hat = x_q * scale
return x_hat, x_q, scale
def fake_quantize(self, x):
"""
伪量化 (训练时使用)
前向传播: 使用量化值
反向传播: 使用直通估计器 (STE)
"""
scale = x.abs().max() / self.qmax
# 前向: 量化 + 反量化
x_q = torch.clamp(torch.round(x / scale), self.qmin, self.qmax)
x_hat = x_q * scale
# 反向: 直通估计器 (STE)
# 梯度直接传过量化操作
x_hat = x + (x_hat - x).detach()
return x_hat
class AsymmetricQuantizer:
"""
非对称量化器
理论:
scale = (max(x) - min(x)) / (2^b - 1)
zero_point = round(-min(x) / scale)
x_q = round(x / scale) + zero_point
x̂ = (x_q - zero_point) × scale
特点:
可以处理非对称分布
更灵活
"""
def __init__(self, bits=8):
self.bits = bits
self.qmin = 0
self.qmax = 2 ** bits - 1
def quantize(self, x):
min_val = x.min()
max_val = x.max()
scale = (max_val - min_val) / (self.qmax - self.qmin)
zero_point = torch.round(-min_val / scale)
zero_point = torch.clamp(zero_point, self.qmin, self.qmax)
x_q = torch.clamp(torch.round(x / scale) + zero_point, self.qmin, self.qmax)
x_hat = (x_q - zero_point) * scale
return x_hat, x_q, scale, zero_point
5.2 训练后量化（PTQ）
训练后量化 (Post-Training Quantization, PTQ):
在模型训练完成后进行量化
不需要重新训练
方法:
1. 动态量化:
运行时计算 scale 和 zero_point
简单，但有额外开销
2. 静态量化:
使用校准数据预先计算 scale
推理时直接使用
3. 混合精度量化:
不同层使用不同精度
敏感层保持高精度
校准 (Calibration):
使用少量数据 (100-1000 个样本)
统计每层的激活值范围
计算最优的 scale 和 zero_point
class PostTrainingQuantization:
"""
训练后量化
理论:
不需要重新训练
使用校准数据确定量化参数
流程:
1. 收集校准数据
2. 前向传播，统计激活值范围
3. 计算 scale 和 zero_point
4. 应用量化
"""
def __init__(self, model, bits=8):
self.model = model
self.bits = bits
self.calibration_data = None
def calibrate(self, calibration_loader):
"""
校准: 收集激活值统计信息
理论:
使用少量数据估计每层的值域
确定量化的 scale 和 zero_point
"""
activation_stats = {}
# 注册 hook 收集激活值
hooks = []
for name, module in self.model.named_modules():
if isinstance(module, (nn.Conv2d, nn.Linear)):
hook = module.register_forward_hook(
self._create_hook(name, activation_stats)
)
hooks.append(hook)
# 前向传播收集统计
self.model.eval()
with torch.no_grad():
for batch in calibration_loader:
self.model(batch)
# 移除 hook
for hook in hooks:
hook.remove()
return activation_stats
def _create_hook(self, name, stats):
def hook_fn(module, input, output):
if name not in stats:
stats[name] = {'min': float('inf'), 'max': float('-inf')}
stats[name]['min'] = min(stats[name]['min'], output.min().item())
stats[name]['max'] = max(stats[name]['max'], output.max().item())
return hook_fn
"""
PTQ 的局限性:
1. 精度下降:
INT8 PTQ 通常有 1-2% 精度损失
INT4 PTQ 可能有 5-10% 精度损失
2. 校准数据依赖:
需要代表性的校准数据
数据分布偏移会导致量化误差
3. 对异常值敏感:
极端值会影响 scale 计算
导致正常值的量化误差增大
"""
5.3 量化感知训练（QAT）
量化感知训练 (Quantization-Aware Training, QAT):
在训练过程中模拟量化效果
让模型学习适应量化误差
核心技术:
1. 伪量化 (Fake Quantization):
前向传播时模拟量化效果
反向传播时使用直通估计器 (STE)
2. 直通估计器 (Straight-Through Estimator):
前向: quantize(x)
反向: gradient 直接传过 (忽略量化不可导)
∂L/∂x ≈ ∂L/∂x̂ (近似)
3. 学习量化参数:
让 scale 和 zero_point 可学习
class QuantizationAwareTraining:
"""
量化感知训练
理论:
在训练时模拟量化效果
让模型学习适应量化误差
关键技术:
1. 伪量化: 前向时模拟量化
2. STE: 反向时梯度直通
3. 学习量化参数: scale 可学习
"""
def __init__(self, model, bits=8):
self.model = model
self.bits = bits
# 替换层为量化版本
self._prepare_model()
def _prepare_model(self):
"""将模型层替换为量化感知版本"""
for name, module in self.model.named_modules():
if isinstance(module, nn.Conv2d):
# 替换为量化卷积
quantized = QuantizedConv2d(module, self.bits)
self._replace_module(name, quantized)
def train_step(self, batch, labels):
"""
训练步骤
前向传播使用伪量化
反向传播使用 STE
"""
# 前向 (包含伪量化)
outputs = self.model(batch)
# 损失
loss = F.cross_entropy(outputs, labels)
# 反向 (STE 自动处理梯度)
loss.backward()
return loss
class StraightThroughEstimator(torch.autograd.Function):
"""
直通估计器 (STE)
理论:
量化操作不可导 (round 操作)
STE 在反向传播时忽略量化
直接将梯度传过
前向: x̂ = quantize(x)
反向: ∂L/∂x = ∂L/∂x̂
这是一个近似，但实践中效果好
"""
@staticmethod
def forward(ctx, x, scale, qmin, qmax):
x_q = torch.clamp(torch.round(x / scale), qmin, qmax)
x_hat = x_q * scale
return x_hat
@staticmethod
def backward(ctx, grad_output):
# 梯度直接传过 (忽略量化)
return grad_output, None, None, None
5.4 混合精度量化
混合精度量化 (Mixed-Precision Quantization):
理论:
不同层对量化的敏感度不同
敏感层使用高精度
不敏感层使用低精度
敏感度分析:
敏感度 = 量化前后的性能差异
高敏感层: 第一层、最后一层、注意力层
低敏感层: 中间的卷积层、FFN
方法:
1. 基于规则: 手动指定精度
2. 基于搜索: 自动搜索最优精度配置
3. 基于敏感度: 根据敏感度分配精度
class MixedPrecisionQuantizer:
"""
混合精度量化
理论:
不同层使用不同的量化精度
敏感层: 高精度 (FP16/INT8)
不敏感层: 低精度 (INT4/INT2)
"""
def __init__(self, model):
self.model = model
self.sensitivity = {}
def analyze_sensitivity(self, calibration_loader):
"""
分析各层的量化敏感度
方法: 逐层量化，测量性能变化
"""
for name, module in self.model.named_modules():
if isinstance(module, (nn.Conv2d, nn.Linear)):
# 量化该层
original_output = self._get_layer_output(name, calibration_loader)
quantized_output = self._quantize_and_get_output(name, calibration_loader)
# 计算敏感度
sensitivity = F.mse_loss(original_output, quantized_output)
self.sensitivity[name] = sensitivity.item()
return self.sensitivity
def assign_precision(self, total_bits_budget):
"""
分配精度
理论:
在总比特预算约束下
最小化量化误差
"""
# 按敏感度排序
sorted_layers = sorted(
self.sensitivity.items(),
key=lambda x: x[1],
reverse=True
)
precision_map = {}
remaining_bits = total_bits_budget
for name, sens in sorted_layers:
if sens > 0.1: # 高敏感
precision_map[name] = 8 # INT8
remaining_bits -= 8
elif sens > 0.01: # 中等敏感
precision_map[name] = 4 # INT4
remaining_bits -= 4
else: # 低敏感
precision_map[name] = 2 # INT2
remaining_bits -= 2
return precision_map
6. 低秩分解
6.1 理论基础
低秩分解 (Low-Rank Factorization):
将大矩阵分解为小矩阵的乘积
原始矩阵: W ∈ ℝ^{m×n}
分解: W ≈ A × B, 其中 A ∈ ℝ^{m×r}, B ∈ ℝ^{r×n}, r << min(m,n)
参数量: m×n → m×r + r×n = r×(m+n)
当 r << min(m,n) 时，大幅减少参数
理论基础:
奇异值分解 (SVD):
W = UΣV^T
其中:
U ∈ ℝ^{m×m}: 左奇异向量
Σ ∈ ℝ^{m×n}: 奇异值矩阵
V ∈ ℝ^{n×n}: 右奇异向量
低秩近似:
只保留最大的 r 个奇异值
W ≈ U_r × Σ_r × V_r^T
近似误差:
‖W - W_r‖_F = √(Σ_{i>r} σ_i²)
Eckart-Young 定理: SVD 给出最优的低秩近似
class LowRankDecomposition:
"""
低秩分解
理论:
将大矩阵分解为小矩阵的乘积
利用矩阵的低秩结构
方法:
1. SVD 分解
2. Tucker 分解
3. CP 分解
"""
@staticmethod
def svd_decomposition(weight, rank):
"""
SVD 分解
W ≈ U_r × Σ_r × V_r^T
理论:
Eckart-Young 定理保证最优性
保留最大的 r 个奇异值
"""
U, S, V = torch.svd(weight)
# 保留前 r 个
U_r = U[:, :rank]
S_r = S[:rank]
V_r = V[:, :rank]
# 重构
W_approx = U_r @ torch.diag(S_r) @ V_r.T
return U_r, S_r, V_r, W_approx
@staticmethod
def choose_rank(weight, energy_threshold=0.99):
"""
选择秩
理论:
保留足够的奇异值以捕获指定比例的能量
能量比 = Σ_{i≤r} σ_i² / Σ_i σ_i²
"""
_, S, _ = torch.svd(weight)
# 计算累积能量
total_energy = (S ** 2).sum()
cumulative_energy = torch.cumsum(S ** 2, dim=0)
# 找到满足阈值的最小秩
rank = torch.searchsorted(cumulative_energy / total_energy, energy_threshold).item()
return rank
"""
SVD 分解的压缩效果:
原始: W ∈ ℝ^{m×n}, 参数量 = m×n
分解: A ∈ ℝ^{m×r}, B ∈ ℝ^{r×n}
参数量 = m×r + r×n = r×(m+n)
压缩比 = m×n / (r×(m+n))
例: m=n=1024, r=64
原始: 1024² = 1,048,576
分解: 64×(1024+1024) = 131,072
压缩比: 8x
"""
6.2 低秩近似在 LLM 中的应用
class LowRankLinear(nn.Module):
"""
低秩线性层
将 W ∈ ℝ^{m×n} 分解为 A ∈ ℝ^{m×r}, B ∈ ℝ^{r×n}
理论:
当 r << min(m,n) 时
大幅减少参数和计算量
与 LoRA 的区别:
LoRA: 冻结原始权重，添加低秩增量
低秩分解: 直接替换原始权重
"""
def __init__(self, m, n, rank):
super().__init__()
self.A = nn.Parameter(torch.randn(m, rank) * 0.01)
self.B = nn.Parameter(torch.randn(rank, n) * 0.01)
self.rank = rank
def forward(self, x):
# x @ (A @ B) = (x @ A) @ B
return (x @ self.A) @ self.B
@staticmethod
def from_linear(linear_layer, rank):
"""从标准线性层创建低秩版本"""
W = linear_layer.weight.data
# SVD 分解
U, S, V = torch.svd(W)
# 保留前 rank 个
A = U[:, :rank] @ torch.diag(S[:rank].sqrt())
B = S[:rank].sqrt().unsqueeze(1) * V[:, :rank].T
low_rank = LowRankLinear(W.shape[0], W.shape[1], rank)
low_rank.A.data = A
low_rank.B.data = B
return low_rank
"""
低秩分解在 Transformer 中的应用:
可以分解的层:
1. 注意力层的 QKV 投影
2. FFN 的权重矩阵
3. Embedding 矩阵
典型压缩:
原始: 4096 × 4096 = 16M 参数
分解: 4096 × 256 + 256 × 4096 = 2M 参数
压缩比: 8x
"""
7. 神经架构搜索
7.1 NAS 概述
神经架构搜索 (Neural Architecture Search, NAS):
自动搜索最优的网络架构
搜索空间:
- 操作类型: 卷积、池化、跳跃连接等
- 连接方式: 串行、并行、密集连接等
- 超参数: 通道数、层数、核大小等
搜索策略:
- 强化学习: 将架构选择视为决策过程
- 进化算法: 模拟自然选择
- 梯度方法: 将离散选择松弛为连续优化
评估策略:
- 完整训练: 每个候选架构都完整训练 (代价高)
- 权重共享: 所有架构共享权重 (代价低)
- 超网络: 训练一个包含所有架构的大网络
7.2 DARTS
论文: "DARTS: Differentiable Architecture Search" (Liu et al., 2019)
核心思想:
将离散的架构选择松弛为连续优化
可以使用梯度下降搜索
方法:
定义混合操作:
o̅(x) = Σ_o α_o · o(x)
其中:
α_o: 架构参数 (可学习)
o(x): 具体操作 (如 3x3 卷积)
使用 softmax 归一化:
α̅ = softmax(α)
训练:
交替优化权重 w 和架构参数 α
min_w,α L_train(w, α)
搜索完成后:
选择 α 最大的操作作为最终架构
class DARTSCell(nn.Module):
"""
DARTS 搜索单元
理论:
将离散选择松弛为连续混合
通过梯度下降学习最优架构
"""
def __init__(self, C, num_nodes=4, num_ops=8):
super().__init__()
self.num_nodes = num_nodes
# 操作列表
self.ops = nn.ModuleList([
nn.ModuleList([
self._build_op(C, op_type)
for op_type in range(num_ops)
])
for _ in range(num_nodes * (num_nodes + 1) // 2)
])
# 架构参数
self.alphas = nn.Parameter(
torch.randn(num_nodes * (num_nodes + 1) // 2, num_ops) * 0.001
)
def forward(self, x):
states = [x]
for i in range(self.num_nodes):
# 混合所有输入
node_input = 0
for j in range(len(states)):
edge_idx = self._edge_index(j, i)
# softmax 权重
weights = F.softmax(self.alphas[edge_idx], dim=-1)
# 混合操作
mixed = sum(w * op(states[j])
for w, op in zip(weights, self.ops[edge_idx]))
node_input += mixed
states.append(node_input)
return states[-1]
def _build_op(self, C, op_type):
"""构建操作"""
ops = [
nn.Conv2d(C, C, 3, padding=1), # 3x3 卷积
nn.Conv2d(C, C, 5, padding=2), # 5x5 卷积
nn.Conv2d(C, C, 3, padding=1, groups=C), # 深度可分离
nn.MaxPool2d(3, stride=1, padding=1), # 最大池化
nn.AvgPool2d(3, stride=1, padding=1), # 平均池化
nn.Identity(), # 跳跃连接
nn.Conv2d(C, C, 1), # 1x1 卷积
nn.Sequential( # 3x3 + 3x3
nn.Conv2d(C, C, 3, padding=1),
nn.ReLU(),
nn.Conv2d(C, C, 3, padding=1)
)
]
return ops[op_type]
8. LLM 专用压缩技术
8.1 LLM 压缩的特殊挑战
LLM 压缩的特殊挑战:
1. 规模巨大:
7B-70B 参数，传统方法不适用
2. 涌现能力:
压缩可能破坏涌现能力
3. 生成质量:
生成任务对误差更敏感
4. 长序列:
KV Cache 占用大量内存
5. 自回归:
每步推理都依赖上一步
LLM 专用方法:
1. GPTQ: 高效的 4-bit 量化
2. AWQ: 激活感知的权重量化
3. SpQR: 稀疏+量化混合
4. LLAMA.cpp: CPU 推理优化
5. vLLM: 高效推理引擎
8.2 GPTQ
论文: "GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers"
(Frantar et al., 2023)
核心思想:
逐列量化权重
使用 Hessian 信息补偿量化误差
理论:
量化误差最小化:
min ‖WX - Q(W)X‖²
其中:
W: 原始权重
Q(W): 量化后的权重
X: 校准数据的激活
使用 Hessian 信息:
H = 2X·X^T
逐列处理:
量化一列后，调整后续列以补偿误差
class GPTQ:
"""
GPTQ 量化器
理论:
逐列量化权重
使用 Hessian 信息补偿误差
一次前向传播即可完成量化
"""
def __init__(self, model, bits=4, group_size=128):
self.model = model
self.bits = bits
self.group_size = group_size
def quantize_layer(self, weight, hessian):
"""
量化单层
理论:
逐列处理
量化一列后，调整后续列补偿误差
"""
rows, cols = weight.shape
Q = torch.zeros_like(weight)
# Hessian 逆
H_inv = torch.linalg.inv(hessian)
for col in range(cols):
# 当前列
w_col = weight[:, col]
# 量化
q_col = self._quantize_column(w_col)
Q[:, col] = q_col
# 计算误差
error = w_col - q_col
# 补偿后续列
if col < cols - 1:
compensation = error.unsqueeze(1) * H_inv[col, col+1:].unsqueeze(0)
weight[:, col+1:] += compensation
return Q
def _quantize_column(self, column):
"""量化单列"""
scale = column.abs().max() / (2 ** (self.bits - 1) - 1)
q = torch.round(column / scale)
q = torch.clamp(q, -(2 ** (self.bits - 1)), 2 ** (self.bits - 1) - 1)
return q * scale
"""
GPTQ 的优势:
1. 高效: 一次前向传播完成量化
2. 准确: 使用 Hessian 补偿，精度损失小
3. 通用: 适用于各种 Transformer 模型
4. 高压缩: 4-bit 量化，4x 压缩
LLaMA-7B:
FP16: 14 GB → INT4: 3.5 GB
精度损失: <1%
"""
8.3 AWQ
论文: "AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration"
(Lin et al., 2024)
核心思想:
不是所有权重同等重要
激活值大的通道对应的权重更重要
理论:
观察: 少数权重通道 (1%) 对应大的激活值
这些通道对量化更敏感
方法: 对重要通道使用更高精度
重要性 = E[|activation|] × |weight|
class AWQ:
"""
AWQ (Activation-aware Weight Quantization)
理论:
激活感知的权重量化
重要通道使用更高精度
优势:
- 不需要反向传播
- 比 GPTQ 更快
- 精度相当
"""
def __init__(self, model, bits=4, group_size=128):
self.model = model
self.bits = bits
self.group_size = group_size
def find_salient_channels(self, weight, activations):
"""
找到重要通道
重要性 = E[|activation|]
"""
# 计算激活的平均幅度
activation_importance = activations.abs().mean(dim=0)
# 找到 top-k 重要通道
k = int(len(activation_importance) * 0.01) # 1%
_, salient_indices = torch.topk(activation_importance, k)
return salient_indices
def quantize_with_scaling(self, weight, salient_indices):
"""
带缩放的量化
对重要通道乘以缩放因子
减小量化误差
"""
# 计算缩放因子
scale = torch.ones(weight.shape[0])
scale[salient_indices] = 2.0 # 重要通道放大
# 应用缩放
weight_scaled = weight * scale.unsqueeze(1)
# 量化
q_weight = self._quantize(weight_scaled)
# 反缩放
q_weight = q_weight / scale.unsqueeze(1)
return q_weight
8.4 结构化剪枝在 LLM 中的应用
class LLMPruner:
"""
LLM 结构化剪枝
可以剪枝:
1. 注意力头
2. FFN 神经元
3. 整个层
"""
def __init__(self, model):
self.model = model
def prune_attention_heads(self, head_importance, keep_ratio=0.8):
"""
剪枝注意力头
理论:
不同注意力头学习不同的模式
有些头可能冗余
重要性:
- 注意力熵: 熵低的头更专注
- 输出幅度: 幅度大的头更重要
- 梯度: 梯度大的头更重要
"""
n_heads = self.model.config.num_attention_heads
n_keep = int(n_heads * keep_ratio)
# 选择重要头
_, important_heads = torch.topk(head_importance, n_keep)
# 创建剪枝后的模型
pruned_model = self._prune_heads(important_heads)
return pruned_model
def prune_ffn_neurons(self, neuron_importance, keep_ratio=0.8):
"""
剪枝 FFN 神经元
理论:
FFN 的中间维度可以剪枝
不重要的神经元对输出影响小
"""
d_ff = self.model.config.intermediate_size
d_keep = int(d_ff * keep_ratio)
# 选择重要神经元
_, important_neurons = torch.topk(neuron_importance, d_keep)
# 创建剪枝后的模型
pruned_model = self._prune_ffn(important_neurons)
return pruned_model
def prune_layers(self, layer_importance, keep_ratio=0.8):
"""
剪枝整个层
理论:
有些层可能冗余
移除整层可以大幅减少计算
挑战:
层间依赖强
需要谨慎选择
"""
n_layers = self.model.config.num_hidden_layers
n_keep = int(n_layers * keep_ratio)
# 选择重要层
_, important_layers = torch.topk(layer_importance, n_keep)
important_layers = important_layers.sort().values
# 创建剪枝后的模型
pruned_model = self._prune_layers(important_layers)
return pruned_model
"""
LLM 层剪枝的理论:
层冗余假设:
深层 Transformer 的某些层可能是冗余的
相邻层的表示可能相似
证据:
- 层间表示相似度高
- 移除某些层后性能下降小
- 跳跃连接可以补偿
方法:
1. 计算层间相似度
2. 移除相似度高的层
3. 微调恢复性能
"""
9. 部署与工程实践
9.1 推理优化框架
┌─────────────────────────────────────────────────────────────────────┐
│ 推理优化框架 │
├───────────────┬─────────────────────────────────────────────────────┤
│ 框架 │ 特点 │
├───────────────┼─────────────────────────────────────────────────────┤
│ TensorRT │ NVIDIA GPU 优化 │
│ │ 自动融合、量化、优化 │
├───────────────┼─────────────────────────────────────────────────────┤
│ ONNX Runtime │ 跨平台推理 │
│ │ 多种硬件支持 │
├───────────────┼─────────────────────────────────────────────────────┤
│ vLLM │ LLM 专用推理引擎 │
│ │ PagedAttention, 连续批处理 │
├───────────────┼─────────────────────────────────────────────────────┤
│ llama.cpp │ CPU 推理优化 │
│ │ GGUF 格式, 量化支持 │
├───────────────┼─────────────────────────────────────────────────────┤
│ TGI │ Hugging Face 推理服务 │
│ │ 生产级部署 │
├───────────────┼─────────────────────────────────────────────────────┤
│ Triton │ NVIDIA 推理服务器 │
│ │ 企业级部署 │
└───────────────┴─────────────────────────────────────────────────────┘
9.2 量化部署
class QuantizedModelDeployment:
"""
量化模型部署
流程:
1. 选择量化方法 (GPTQ/AWQ/动态量化)
2. 量化模型
3. 选择推理框架
4. 优化和部署
"""
@staticmethod
def export_to_onnx(model, input_shape, save_path):
"""导出 ONNX 格式"""
dummy_input = torch.randn(input_shape)
torch.onnx.export(
model,
dummy_input,
save_path,
opset_version=13,
input_names=['input'],
output_names=['output'],
dynamic_axes={
'input': {0: 'batch', 1: 'seq_len'},
'output': {0: 'batch', 1: 'seq_len'}
}
)
@staticmethod
def optimize_tensorrt(onnx_path, engine_path, fp16=True, int8=False):
"""TensorRT 优化"""
import tensorrt as trt
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.max_workspace_size = 1 << 30
if fp16:
config.set_flag(trt.BuilderFlag.FP16)
if int8:
config.set_flag(trt.BuilderFlag.INT8)
engine = builder.build_engine(network, config)
with open(engine_path, 'wb') as f:
f.write(engine.serialize())
"""
部署决策树:
GPU 推理:
低延迟 → TensorRT (FP16/INT8)
高吞吐 → vLLM (连续批处理)
CPU 推理:
llama.cpp (GGUF INT4)
ONNX Runtime (INT8)
边缘设备:
TFLite (移动端)
CoreML (Apple)
ONNX Runtime Mobile
"""
10. 理论分析与前沿
10.1 压缩理论
信息论视角:
模型压缩可以看作信息压缩
原始模型的信息: I(W) = -Σ p(w) log p(w)
压缩后: I(Q(W)) < I(W)
率失真理论 (Rate-Distortion):
给定失真上限 D，最小化率 R (存储空间)
R(D) = min I(Q(W)) s.t. E[d(W, Q(W))] ≤ D
含义:
存在理论下限
在给定精度损失下，有最小的压缩比
泛化理论:
压缩可能改善泛化
最小描述长度 (MDL) 原理:
最好的模型是能用最短描述长度编码数据的模型
PAC-Bayes 界:
泛化误差 ≤ 复杂度 + 拟合误差
压缩减少复杂度 → 可能改善泛化
10.2 前沿研究方向
模型压缩的前沿:
1. 极端压缩:
1-bit LLM (BitNet)
2-bit 量化
理论极限在哪里?
2. 无损压缩:
能否在不损失性能的情况下压缩?
知识蒸馏 + 量化 + 剪枝的组合
3. 自适应压缩:
根据输入动态调整模型大小
简单输入用小模型，复杂输入用大模型
4. 压缩与训练的结合:
从头训练压缩模型
而非压缩预训练模型
5. 硬件协同设计:
为特定压缩方法设计硬件
如: 2:4 稀疏性的 GPU 支持
6. 多模态压缩:
视觉-语言模型的压缩
不同模态的压缩比
附录
A. 压缩方法对比
┌─────────────────────────────────────────────────────────────────────┐
│ 压缩方法对比 │
├───────────────┬────────┬────────┬────────┬──────────────────────────┤
│ 方法 │ 压缩比 │ 精度损失│ 加速比 │ 特点 │
├───────────────┼────────┼────────┼────────┼──────────────────────────┤
│ 知识蒸馏 │ 2-10x │ 1-3% │ 2-10x │ 需要训练 │
│ 非结构化剪枝 │ 5-20x │ 1-5% │ 1-2x │ 需要稀疏硬件 │
│ 结构化剪枝 │ 1.5-3x │ 2-5% │ 1.5-3x │ 直接加速 │
│ INT8 量化 │ 4x │ 0.5-2% │ 2-4x │ 简单，广泛支持 │
│ INT4 量化 │ 8x │ 1-5% │ 2-4x │ LLM 常用 │
│ 低秩分解 │ 2-8x │ 1-3% │ 2-8x │ 理论优美 │
│ NAS │ - │ 0-1% │ - │ 自动设计 │
└───────────────┴────────┴────────┴────────┴──────────────────────────┘
B. 核心公式速查
| 公式 | 含义 |
|---|---|
| L = α·T²·KL(p_T‖p_S) + (1-α)·CE(y, p_S) | 知识蒸馏损失 |
| x_q = round(x/scale) | 量化公式 |
| W ≈ U_r·Σ_r·V_r^T | SVD 低秩分解 |
| α̅ = softmax(α) | DARTS 架构参数 |
| ‖W - Q(W)X‖² | 量化误差 |
C. 推荐资源
- Hinton, G., et al. (2015). Distilling the Knowledge in a Neural Network
- Han, S., et al. (2015). Learning both Weights and Connections for Efficient Neural Networks
- Frantar, E., et al. (2023). GPTQ: Accurate Post-Training Quantization
- Lin, J., et al. (2024). AWQ: Activation-aware Weight Quantization
- Liu, H., et al. (2019). DARTS: Differentiable Architecture Search
免费领 150 小时云算力，进群参与显卡、AI PC 幸运抽奖
更多推荐
0
0- 0
已为社区贡献3条内容
所有评论(0)