模型轻量化是深度学习部署中的关键技术,旨在在保持模型性能(如准确率)的前提下,显著减少模型的参数量、计算量(FLOPs)、内存占用和推理延迟,使其适用于资源受限的设备(如移动端、嵌入式设备、IoT设备等)。

通过系统化实施,可在移动端实现 <10ms 延迟、<5MB 模型体积,同时保持 90%+ 原始精度。

  • 精度-效率权衡:轻量化必然带来精度损失,需根据业务容忍度调整。

  • 硬件适配:INT8在GPU/NPU上加速明显,但在CPU上可能收效甚微。

  • 端到端延迟:不仅看FLOPs,还要考虑内存带宽、缓存命中率。

  • 动态输入:部分轻量化方法(如通道剪枝)不支持动态shape。

(内容逐步完善中)

技术手段与路线

网络结构设计

从源头设计轻量级网络结构。

  • MobileNet(V1/V2/V3):使用深度可分离卷积(Depthwise Separable Convolution)

  • ShuffleNet(V1/V2):使用通道混洗(Channel Shuffle)和分组卷积

  • EfficientNet:通过复合缩放(Compound Scaling)平衡深度、宽度和分辨率

  • GhostNet:通过廉价操作生成“幻影”特征图

实施步骤

  1. 需求分析:明确目标设备(如手机、嵌入式芯片)、延迟/功耗/精度要求。

  2. 选择基线架构:根据任务(图像分类、目标检测等)选择合适的轻量网络(如MobileNetV2用于分类)。

  3. 调整超参数:调节宽度乘子(width multiplier)、分辨率等控制模型大小。

  4. 训练与验证:在目标任务数据集上训练,并验证精度与速度的平衡。

模型剪枝(Pruning)

模型剪枝是一种模型压缩技术,通过移除神经网络中不重要的参数(权重、神经元、通道等)来减少模型大小和计算量,同时尽量保持模型性能。

移除冗余的权重、通道或层。

  • 权重剪枝(Unstructured Pruning):移除绝对值小的权重,形成稀疏矩阵。需要稀疏计算支持(如专用硬件或库)。
  • 通道/结构化剪枝(Structured Pruning):移除整个卷积核或通道,保持稠密结构,兼容通用硬件。

原始模型 → 识别重要参数 → 移除不重要参数 → 微调 → 剪枝后模型

剪枝的粒度级别

  • 权重级”: “移除单个权重参数”,

  • 神经元级”: “移除整个神经元”,

  • 通道级”: “移除整个特征通道”,

  • 层级”: “移除整个网络层”,

  • 块级”: “移除网络块(如ResNet块)”

剪枝技术和类别

  • 结构剪枝

    • 非结构化剪枝 - 移除单个权重

    • 结构化剪枝 - 移除整个通道/神经元

  • 剪枝技术

    • 迭代剪枝 - 逐步剪枝并微调
    • 基于正则化的剪枝
      • L1正则化损失 - 促进稀疏性
      • 组Lasso损失 - 促进结构化稀疏性
      • 组合损失函数

核心算法

基于幅度梯度 算法的剪枝

  • 基于幅度

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    class MagnitudeBasedPruning:
    """基于权重大小的剪枝方法"""

    def __init__(self, pruning_rate=0.2):
    self.pruning_rate = pruning_rate

    def global_magnitude_prune(self, model):
    """全局幅度剪枝"""
    all_weights = []

    # 收集所有权重
    for name, param in model.named_parameters():
    if 'weight' in name and len(param.shape) > 1: # 只处理权重矩阵
    all_weights.append(param.data.abs().view(-1))

    all_weights = torch.cat(all_weights)

    # 计算全局阈值
    threshold = torch.quantile(all_weights, self.pruning_rate)

    # 应用剪枝
    masks = {}
    for name, param in model.named_parameters():
    if 'weight' in name and len(param.shape) > 1:
    mask = param.data.abs() > threshold
    masks[name] = mask
    param.data *= mask.float()

    return masks

    def layer_wise_magnitude_prune(self, model):
    """逐层幅度剪枝"""
    masks = {}

    for name, module in model.named_modules():
    if isinstance(module, (nn.Linear, nn.Conv2d)):
    # 计算该层的阈值
    weights = module.weight.data.abs().view(-1)
    threshold = torch.quantile(weights, self.pruning_rate)

    # 创建掩码
    mask = module.weight.data.abs() > threshold
    masks[name] = mask

    # 应用剪枝
    module.weight.data *= mask.float()

    return masks
  • 基于梯度

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    class GradientBasedPruning:
    """基于梯度信息的剪枝方法"""

    def __init__(self, pruning_rate=0.2):
    self.pruning_rate = pruning_rate

    def compute_weight_importance(self, model, dataloader, criterion):
    """计算权重重要性(基于梯度)"""
    model.train()
    importance_scores = {}

    # 初始化重要性分数
    for name, param in model.named_parameters():
    if 'weight' in name:
    importance_scores[name] = torch.zeros_like(param.data)

    # 计算梯度并累积重要性
    for batch_idx, (data, target) in enumerate(dataloader):
    model.zero_grad()
    output = model(data)
    loss = criterion(output, target)
    loss.backward()

    # 累积梯度信息
    for name, param in model.named_parameters():
    if 'weight' in name and param.grad is not None:
    # 使用梯度×权重作为重要性指标
    importance_scores[name] += (param.data * param.grad).abs()

    return importance_scores

    def gradient_based_prune(self, model, dataloader, criterion):
    """基于梯度的剪枝"""
    importance_scores = self.compute_weight_importance(model, dataloader, criterion)

    # 计算全局阈值
    all_scores = torch.cat([score.view(-1) for score in importance_scores.values()])
    threshold = torch.quantile(all_scores, self.pruning_rate)

    # 应用剪枝
    masks = {}
    for name, param in model.named_parameters():
    if 'weight' in name:
    mask = importance_scores[name] > threshold
    masks[name] = mask
    param.data *= mask.float()

    return masks

实施步骤

  1. 预训练模型:在目标任务上训练一个完整模型。

  2. 重要性评估:

    • L1/L2范数:通道权重的L1范数越小越不重要。
    • 梯度信息:使用泰勒展开估计移除某通道对损失的影响。
    • 注意力机制:如使用BN层的γ系数作为通道重要性指标。
  3. 剪枝策略:

    • 一次性剪枝(One-shot):直接剪掉一定比例。
    • 迭代剪枝(Iterative):逐步剪枝 + 微调,效果更好。
  4. 微调(Fine-tuning):恢复因剪枝导致的精度下降。

  5. 导出剪枝后模型:移除冗余参数,生成紧凑模型。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
class AutoPruner:
"""自动剪枝框架"""

def __init__(self, model, dataloader, criterion):
self.model = model
self.dataloader = dataloader
self.criterion = criterion

def sensitivity_analysis(self):
"""敏感性分析 - 确定各层对剪枝的敏感度"""
sensitivity_scores = {}
original_accuracy = self.evaluate_model()

for name, module in self.model.named_modules():
if isinstance(module, (nn.Linear, nn.Conv2d)):
# 测试不同剪枝比例的影响
sensitivities = []

for amount in [0.1, 0.3, 0.5, 0.7]:
# 创建临时副本
temp_model = copy.deepcopy(self.model)
temp_module = dict(temp_model.named_modules())[name]

# 应用剪枝
prune.l1_unstructured(temp_module, name='weight', amount=amount)

# 评估性能
accuracy = self.evaluate_model(temp_model)
accuracy_drop = original_accuracy - accuracy
sensitivities.append(accuracy_drop)

sensitivity_scores[name] = {
'sensitivities': sensitivities,
'avg_drop': sum(sensitivities) / len(sensitivities)
}

return sensitivity_scores

def evaluate_model(self, model=None):
"""评估模型性能"""
if model is None:
model = self.model

model.eval()
correct = 0
total = 0

with torch.no_grad():
for data, target in self.dataloader:
output = model(data)
pred = output.argmax(dim=1, keepdim=True)
correct += pred.eq(target.view_as(pred)).sum().item()
total += target.size(0)

return correct / total

def auto_prune(self, target_sparsity=0.7):
"""自动剪枝"""
# 敏感性分析
sensitivity_scores = self.sensitivity_analysis()

# 根据敏感度分配剪枝比例
pruning_plan = self.create_pruning_plan(sensitivity_scores, target_sparsity)

# 应用剪枝
pruner = ComprehensivePruner(self.model)
pruner.apply_pruning(pruning_plan)

return pruning_plan

def create_pruning_plan(self, sensitivity_scores, target_sparsity):
"""创建剪枝计划"""
pruning_plan = {}

# 按敏感度排序(敏感度低的层可以剪枝更多)
sorted_layers = sorted(sensitivity_scores.items(),
key=lambda x: x[1]['avg_drop'])

total_layers = len(sorted_layers)

for i, (layer_name, scores) in enumerate(sorted_layers):
# 敏感度低的层分配更高的剪枝比例
base_amount = target_sparsity * (1 - i / total_layers)

# 根据具体敏感度调整
sensitivity_factor = 1 - min(scores['avg_drop'] * 10, 0.8)
final_amount = base_amount * sensitivity_factor

pruning_plan[layer_name] = {
'method': 'l1_unstructured',
'amount': final_amount
}

return pruning_plan

量化感知(Quantization)

将浮点数(FP32)转换为低比特整数(INT8/INT4)或浮点(FP16)。需校准确定量化范围(min/max或scale/zero-point)

  • 训练后量化(PTQ, Post-Training Quantization):无需重新训练,速度快。

  • 量化感知训练(QAT, Quantization-Aware Training):在训练中模拟量化误差,精度更高。

指标 FP32模型 INT8量化后 INT4量化后
模型大小(7B参数) ~28GB ~7GB ~3.5GB
显存占用(推理) 14GB 7GB 3.5GB
推理速度 基准1x 1.5-2x更快 2-3x更快
精度损失 <1% (通常可接受) 1-3% (可能需调优)
  • 原始精度:模型训练时通常使用FP32(32位浮点数),每个参数占用4字节。

  • 量化后:将参数转换为低精度格式(如INT8、INT4),每个参数仅占1字节甚至0.5字节。

    • 例如:将权重从 0.8732(FP32)近似为 0.875(INT8)。

量化的主要类型

类型 说明 典型应用场景
训练后量化 对已训练好的模型直接量化(无需重新训练) 快速部署,资源有限环境
量化感知训练 在训练过程中模拟量化误差,使模型适应低精度 高精度要求的微调场景
动态量化 推理时动态量化权重和激活值 实时性要求高的场景
静态量化 预先校准量化参数(如缩放因子),推理时固定 移动端/嵌入式设备

量化的实现技术

方法路径 核心思路 适用场景
GPTQ量化 一种训练后量化方法,尤其适合降低大语言模型的显存占用和提升推理速度。 希望获得极致性能,对精度损失有一定容忍度。
PyTorch原生量化 使用torch.ao.quantization等PyTorch内置模块,支持动态和静态量化。 需要官方支持,与PyTorch生态紧密结合。
TensorRT集成 通过转换模型至ONNX格式,再利用TensorRT进行优化和量化,能显著提升推理速度。 生产环境部署,追求NVIDIA硬件上的最高推理性能。

GPTQ 量化

PyTorch 量化

TensorRT 集成

具体实施

实施步骤(以QAT为例)

  1. 插入伪量化节点:在训练图中插入FakeQuant操作,模拟量化过程。

  2. 使用量化感知损失函数:保持梯度可传。

  3. 训练/微调模型:通常只需少量epoch。

  4. 转换为实际量化模型:将FakeQuant替换为真实INT8操作。

  5. 部署:使用支持INT8的推理引擎(如TensorRT、ONNX Runtime、TFLite)。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
import torch
import torch.nn as nn
from Quantization.data import get_sample_input
from torch.ao.quantization import (
get_default_qconfig_mapping,
prepare,
convert,
)

def quantize_model_fp32_to_int8(model_fp32: nn.Module, calib_data) -> torch.nn.Module:
"""
使用训练后量化(PTQ)将 FP32 模型转为 INT8(仅 CPU 支持)

Args:
model_fp32: 原始 FP32 模型(必须是 eval 模式)
calib_data: 校准数据,用于确定量化范围(min/max)

Returns:
量化后的 INT8 模型
"""
# Step 1: 设置量化配置(仅支持 CPU 后端)
qconfig_mapping = get_default_qconfig_mapping("x86") # 或 "fbgemm"(Linux)

# Step 2: 准备模型(插入 observer)
example_inputs = (calib_data[0], calib_data[1])
model_prepared = prepare(model_fp32, qconfig_mapping, example_inputs)

# Step 3: 校准(用少量数据统计激活值范围)
print("正在校准模型...")
with torch.no_grad():
for _ in range(10): # 10 个 batch 足够
x, x_lens = get_sample_input()
model_prepared(x, x_lens)

# Step 4: 转换为量化模型
model_int8 = convert(model_prepared)
print("量化完成!")
return model_int8

为什么量化后模型大小没变?

原因:state_dict() 保存的是 量化参数(scale/zero_point)+ 量化权重(int8),但 PyTorch 默认以 float32 格式序列化所有张量!即使权重是 int8,当你调用 torch.save(model.state_dict()) 时:

  • PyTorch 会把 int8 张量自动转换为 float32 存储(为了兼容性);

  • 同时还保存了 scalezero_point 等额外参数;

最终文件大小 ≈ 原始 FP32 模型,甚至更大!

量化的核心收益在运行时,不在存储文件。部署时应使用 TorchScript / ONNX 格式。

知识蒸馏(Distillation)

知识蒸馏(Knowledge Distillation)是一种模型压缩技术,其核心思想是将一个庞大、复杂但性能优异的模型(教师模型)的知识转移到一个更小、更高效的模型(学生模型)中。

用大模型(Teacher)指导小模型(Student)学习。

教师模型 (大而复杂) → 知识转移 → 学生模型 (小而高效)。

知识蒸馏的成功, 80% 取决于训练数据质量 ,且Teacher 和 Student 差距不宜过大

优势

  • 模型压缩:大幅减少参数量和计算量

  • 性能保持:学生模型性能接近教师模型

  • 推理加速:更快的推理速度

  • 部署友好:适合资源受限环境

挑战

  • 教师模型质量依赖:教师模型质量直接影响蒸馏效果

  • 超参数敏感:温度、权重系数等需要仔细调优

  • 训练复杂度:需要同时训练教师和学生模型

  • 知识损失:不可避免会损失部分知识

核心原理

  • 软标签:教师模型输出的概率分布包含更多信息

  • 暗知识:类别之间的关系等隐含知识

  • 温度参数:控制概率分布的平滑程度

软标签与硬标签

1
2
3
4
5
6
7
8
9
10
11
12
13
import torch
import torch.nn as nn
import torch.nn.functional as F

# 硬标签 vs 软标签示例
def demonstrate_labels():
# 硬标签: [0, 0, 1, 0, 0]
hard_labels = torch.tensor([2]) # 只是类别索引

# 软标签: [0.1, 0.2, 0.5, 0.15, 0.05]
soft_labels = torch.tensor([0.1, 0.2, 0.5, 0.15, 0.05])

return hard_labels, soft_labels

温度缩放:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
class KnowledgeDistillationLoss(nn.Module):
def __init__(self, temperature=3.0, alpha=0.7):
super().__init__()
self.temperature = temperature
self.alpha = alpha
self.kl_loss = nn.KLDivLoss(reduction="batchmean")
self.ce_loss = nn.CrossEntropyLoss()

def forward(self, student_logits, teacher_logits, labels):
# 应用温度缩放
student_soft = F.log_softmax(student_logits / self.temperature, dim=1)
teacher_soft = F.softmax(teacher_logits / self.temperature, dim=1)

# 蒸馏损失
distillation_loss = self.kl_loss(student_soft, teacher_soft) * (self.temperature ** 2)

# 学生损失
student_loss = self.ce_loss(student_logits, labels)

# 总损失
total_loss = self.alpha * distillation_loss + (1 - self.alpha) * student_loss
return total_loss

训练策略

渐进式蒸馏

1
2
3
4
5
6
7
8
9
10
11
12
13
class ProgressiveDistillation:
"""渐进式知识蒸馏"""

def __init__(self, total_epochs):
self.total_epochs = total_epochs

def get_distillation_weights(self, current_epoch):
"""动态调整蒸馏权重"""
# 早期更依赖教师,后期更依赖真实标签
alpha = max(0.1, 0.7 * (1 - current_epoch / self.total_epochs))
temperature = max(1.0, 3.0 * (1 - current_epoch / self.total_epochs))

return alpha, temperature

注意力转移

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
class AttentionDistillation:
"""注意力机制蒸馏"""

def __init__(self):
self.mse_loss = nn.MSELoss()

def compute_attention_loss(self, student_attentions, teacher_attentions):
"""计算注意力蒸馏损失"""
attention_loss = 0
num_layers = min(len(student_attentions), len(teacher_attentions))

for i in range(num_layers):
s_attn = student_attentions[i] # [batch, heads, seq_len, seq_len]
t_attn = teacher_attentions[i]

# 确保头数匹配
if s_attn.shape[1] != t_attn.shape[1]:
# 平均池化调整头数
if s_attn.shape[1] < t_attn.shape[1]:
# 学生头数少,对教师注意力求平均
t_attn = t_attn.mean(dim=1, keepdim=True)
else:
# 学生头数多,复制教师注意力
t_attn = t_attn.repeat(1, s_attn.shape[1] // t_attn.shape[1], 1, 1)

layer_loss = self.mse_loss(s_attn, t_attn)
attention_loss += layer_loss

return attention_loss / num_layers

实施步骤

  1. 训练Teacher模型:高精度但复杂的大模型。

  2. 设计Student模型:结构更轻量(如层数更少、通道更窄)。

  3. 定义蒸馏损失:

    • 软标签损失(Soft Target):使用Teacher输出的softmax logits(温度缩放)。
    • 特征图对齐:中间层特征的L2或注意力对齐。
  4. 联合训练:Student同时学习真实标签和Teacher的输出。

  5. 部署Student模型:独立使用,无需Teacher。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
import torch
import torch.nn as nn
from transformers import GPT2Config, GPT2LMHeadModel

class TeacherStudentDistillation:
"""完整的知识蒸馏实现"""

def __init__(self, teacher_model, student_model, temperature=3.0, alpha=0.7):
self.teacher_model = teacher_model
self.student_model = student_model
self.temperature = temperature
self.alpha = alpha

# 冻结教师模型参数
for param in self.teacher_model.parameters():
param.requires_grad = False

# 损失函数
self.kl_loss = nn.KLDivLoss(reduction="batchmean")
self.ce_loss = nn.CrossEntropyLoss()
self.mse_loss = nn.MSELoss()

def create_models():
"""创建教师和学生模型"""
# 教师模型 (大模型)
teacher_config = GPT2Config(
n_layer=12, # 12层
n_head=12, # 12个注意力头
n_embd=768 # 768维嵌入
)
teacher_model = GPT2LMHeadModel(teacher_config)

# 学生模型 (小模型)
student_config = GPT2Config(
n_layer=6, # 6层
n_head=6, # 6个注意力头
n_embd=384 # 384维嵌入
)
student_model = GPT2LMHeadModel(student_config)

return teacher_model, student_model

def compute_distillation_loss(self, student_logits, teacher_logits, labels):
"""计算蒸馏损失"""
# 响应式知识蒸馏
student_soft = F.log_softmax(student_logits / self.temperature, dim=-1)
teacher_soft = F.softmax(teacher_logits / self.temperature, dim=-1)

distillation_loss = self.kl_loss(student_soft, teacher_soft) * (self.temperature ** 2)

# 学生任务损失
task_loss = self.ce_loss(student_logits.view(-1, student_logits.size(-1)),
labels.view(-1))

# 组合损失
total_loss = self.alpha * distillation_loss + (1 - self.alpha) * task_loss
return total_loss

def train_step(self, input_ids, attention_mask, labels):
"""训练步骤"""
# 教师模型前向传播
with torch.no_grad():
teacher_outputs = self.teacher_model(
input_ids=input_ids,
attention_mask=attention_mask,
labels=labels
)
teacher_logits = teacher_outputs.logits

# 学生模型前向传播
student_outputs = self.student_model(
input_ids=input_ids,
attention_mask=attention_mask,
labels=labels
)
student_logits = student_outputs.logits

# 计算蒸馏损失
loss = self.compute_distillation_loss(student_logits, teacher_logits, labels)
return loss

Token 循环问题

这是知识蒸馏失败的典型症状,模型陷入重复 token 循环

Student (原始): The future of AI is the real question.
One of the first things I learned about AI is that it’s almost impossible to write algorithms in this way. It’s also a bit difficult to understand how any algorithm can perform any task. For example, in one of

Student (蒸馏后): The future of AI is and and and and and and and(大量重复token,也可能是其他符号或字符)

出现此种问题的可能原因为:

  • 训练数据不足/质量差

    模型无法学习通用语言模式,只能记住训练数据中的模式。如果训练数据中 , 出现频率高,模型会过度生成。

    尽量使用真实训练数据

  • 蒸馏参数不当

    • alpha=0.7 过高:Student 过度依赖 Teacher 的软标签。可以适当增加硬标签权重(0.3~0.5)
    • temperature=2.0 过高:Teacher 分布过于平滑,失去区分度。可以减少平滑(1.0~2.0)
  • 无真实标签监督

    硬标签损失(CE)权重仅 30%,模型未充分学习基本语言建模能力。

    可以适当增加硬标签损失权重

  • Teacher 和 Student 差距过大

    gpt2-medium (355M) → gpt2 (124M),小模型难以模仿大模型的复杂行为。

    可以使用更小的 Teacher

同时可添加惩罚重复属性:

1
2
3
4
5
6
7
8
outputs = model.generate(
**inputs,
max_new_tokens=max_length,
repetition_penalty=1.2, # 惩罚重复 token
do_sample=True,
temperature=0.7,
pad_token_id=tokenizer.eos_token_id
)

低秩分解(Low-Rank)

将一个大权重矩阵/张量近似分解为多个小矩阵/张量的乘积,从而减少参数量和计算量

  • 将卷积核 分解为两个更小的卷积。
    $$
    W∈R
    C
    out​×C
    in​×k×k
    $$

  • 使用SVD分解全连接层。

想象一个复杂的变换需要1000个输入和1000个输出,那么它的权重矩阵 W 的大小是 1000×1000,共有100万个参数。

低秩分解发现,这个变换的内在“自由度”或“信息量”其实没那么高(即它是低秩的)。

  1. 先将1000维输入投影到一个低维空间(比如50维)。这对应一个矩阵 A (1000×50)。

  2. 再从这个50维空间恢复到1000维输出。这对应一个矩阵 B (50×1000)。

于是:W ≈ B × A

参数总量从: 1000 × 1000 = 1,000,000
减少到: 1000 × 50 + 50 × 1000 = 100,000
压缩率高达90%。

实施步骤

  1. 对预训练模型的权重进行SVD或CP/Tucker分解。

  2. 替换原层为分解后的多层结构。

  3. 微调模型恢复精度。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
import torch
import torch.nn as nn

def decompose_linear_svd(linear_layer: nn.Linear, rank_ratio=0.5):
"""
使用 SVD 对 Linear 层进行低秩分解
输入: Linear(in_features=n, out_features=m)
输出: (Linear(n, r), Linear(r, m))
"""
if not isinstance(linear_layer, nn.Linear):
raise TypeError(f"Expected nn.Linear, got {type(linear_layer)}")

W = linear_layer.weight.data # [m, n] = [256, 2560]
m, n = W.shape

max_rank = min(m, n) # = 256
r = int(rank_ratio * max_rank)
r = max(1, min(r, max_rank - 1))

U, S, Vh = torch.linalg.svd(W, full_matrices=False)
# print(f" U shape: {U.shape}") # [256, 256]
# print(f" S shape: {S.shape}") # [256]
# print(f" Vh shape: {Vh.shape}") # [256, 2560] ← 必须是这个!

U_r = U[:, :r] # [256, r]
S_r = S[:r] # [r]
Vh_r = Vh[:r, :] # [r, 2560] ← 关键!

sqrt_S = torch.sqrt(S_r)
B = sqrt_S.unsqueeze(1) * Vh_r # [r, 2560]
A = U_r * sqrt_S.unsqueeze(0) # [256, r]

# print(f" B shape: {B.shape}") # 应为 [r, 2560]
# print(f" A shape: {A.shape}") # 应为 [256, r]

fc1 = nn.Linear(n, r, bias=False)
fc2 = nn.Linear(r, m, bias=linear_layer.bias is not None)

with torch.no_grad():
fc1.weight.copy_(B) # [r, n] = [r, 2560]
fc2.weight.copy_(A) # [m, r] = [256, r]
if linear_layer.bias is not None:
fc2.bias.copy_(linear_layer.bias.data)

return fc1, fc2
  • Linear 层W∈Rm×nUV ,其中 U∈Rm×r,V∈Rr×nr≪min(m,n)

  • Conv2d 层:将卷积核张量分解为多个低秩张量(如 CP 分解、Tucker 分解)

核心参数

关键参数 说明 取值范围 改变的影响
1. 秩(Rank)r 分解后的中间维度 1 \leq r < \min(m, n) 最核心参数!
• r ↑ → 精度 ↑,压缩率 ↓
• r ↓ → 压缩率 ↑,精度 ↓(可能崩溃)
2. 秩比例(Rank Ratio)α r = \alpha \cdot \min(m, n) 0.1 \sim 0.9 • α=0.3:高压缩(~70%),高精度损失
• α=0.7:中压缩(~30%),低精度损失
• α>0.8:几乎无压缩
3. 分解目标层 选择哪些层分解 大矩阵优先 • FFN/Linear 层:高收益(参数多)
• 小卷积层(3×3):可能增参
• Embedding/LM Head:通常不分解
4. 分解方式 SVD / CP / Tucker 等 SVD(主流) • SVD:理论最优,适合 Linear
• CP/Tucker:适合卷积核张量
• NMF:非负约束,适合特定场景
5. 是否微调(Fine-tuning) 分解后是否训练 是 / 否 • 无微调:精度损失大(尤其 α<0.5)
• 有微调:可恢复 80%+ 性能

改变秩 r 会直接引发一个典型的 “精度-效率”权衡。值过低压缩率越好,但是分解后的模型输出大量重复 token

关键指标 r 增大(更接近原始矩阵) r 减小(更激进压缩)
模型精度 ↑ 提升 • 重建误差小,更接近原始模型性能。 • 保留更多任务相关特征。 ↓ 下降 • 重建误差大,信息丢失严重。 • 可能导致模型准确率显著降低。
参数量 ↑ 增加 • 分解后的矩阵更大。 ↓ 减少 • 分解后的矩阵更小,压缩率更高。
计算量(FLOPs) ↑ 增加 • 需要进行更多次矩阵乘法。 ↓ 减少 • 计算量显著降低,加速效果更明显。
内存占用 ↑ 增加 • 需要存储更多的参数。 ↓ 减少 • 内存占用显著降低。
过拟合风险 ↑ 增加 • 模型容量相对较大,在小型数据集上可能过拟合。 ↓ 减少 • 模型容量小,起到正则化作用,可能缓解过拟合。
适用场景 • 对精度要求高的任务。
• 原始模型冗余度较低。
• 极度资源受限的边缘设备。
• 对延迟要求极高的场景。
• 原始模型冗余度极高。

秩比例

rank_ratio r (c_fc) 生成质量 参数压缩率
0.3 230 重复 68%
0.5 384 可用 50%
0.7 537 接近原始 30%

建议的配置策略:

策略 rank_ratio 是否微调 生成质量 压缩率
仅分解 c_proj 0.5 良好 25%
分解全部 + 微调 0.5 接近原始 50%
分解全部(无微调) 0.5 轻微重复 50%
分解全部(无微调) 0.3 严重重复 68%

目标层

层类型 示例 是否推荐分解 原因
FFN 中间层 GPT-2 c_fc (768→3072) 谨慎 需高容量编码语义,r 必须大
FFN 投影层 GPT-2 c_proj (3072→768) 推荐 信息已压缩,可安全分解
Attention QKV W_q, W_k, W_v 谨慎 影响注意力质量,需高秩
分类头 Linear(768→1000) 可分解 任务特定,冗余度高
Embedding Token Embedding 不推荐 稀疏激活,分解收益低

二值化/三值化

二值化(Binarization)和三值化(Ternarization)是极致模型压缩技术,将权重从 FP32 压缩到 1-bit(±1)或 2-bit(-1, 0, +1),适用于超低功耗边缘设备(如 MCU、IoT 传感器)。

二值化/三值化是“空间换精度”的极致压缩

  • 优势:模型 <100KB,功耗极低
  • 代价:精度损失,训练复杂
  • 适用:简单任务 + 超低功耗设备
  1. 二值化(Binary Weight Networks, BWN)

    权重 W∈Rm×nW**b∈{−1,+1}m×n。前向计算:y=WbxαWbxα 为缩放因子)

    压缩率:32x(FP32 → 1-bit)

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    class BinaryQuantize(torch.autograd.Function):
    """
    二值化函数(带 STE)
    前向: sign(x)
    反向: 梯度直通(仅 |x|<=1 时传递)
    """

    @staticmethod
    def forward(ctx, input):
    ctx.save_for_backward(input)
    # 二值化: >0 → +1, <=0 → -1
    out = torch.sign(input)
    # 处理 0(PyTorch sign(0)=0,我们设为 +1)
    out[out == 0] = 1
    return out

    @staticmethod
    def backward(ctx, grad_output):
    input, = ctx.saved_tensors
    # STE: 梯度仅在 |input| <= 1 时传递
    grad_input = grad_output.clone()
    grad_input[torch.abs(input) > 1] = 0
    return grad_input
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    class BinaryConv2d(nn.Module):
    """二值化卷积层"""

    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0):
    super().__init__()
    self.in_channels = in_channels
    self.out_channels = out_channels
    self.kernel_size = kernel_size
    self.stride = stride
    self.padding = padding

    # 浮点权重(用于训练)
    self.weight = nn.Parameter(torch.randn(out_channels, in_channels, kernel_size, kernel_size))
    self.bias = nn.Parameter(torch.zeros(out_channels))

    # 缩放因子 α = mean(|W|)
    self.alpha = nn.Parameter(torch.ones(out_channels, 1, 1, 1))

    def forward(self, x):
    # 1. 计算缩放因子 α
    alpha = self.weight.abs().mean(dim=(1, 2, 3), keepdim=True)

    # 2. 二值化权重
    weight_b = BinaryQuantize.apply(self.weight)

    # 3. 前向计算: y = α * (x ⊗ W_b) + b
    out = F.conv2d(x, weight_b, None, self.stride, self.padding)
    out = out * alpha + self.bias.view(1, -1, 1, 1)
    return out
  2. 三值化(Ternary Weight Networks, TWN)

    权重 W∈Rm×nW**t∈{−1,0,+1}m×n。保留重要权重(非零),移除冗余(零)

    压缩率:16x(FP32 → 2-bit,含稀疏性)

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    class TernaryQuantize(torch.autograd.Function):
    """
    三值化函数(带 STE)
    阈值 δ = 0.7 * mean(|W|)
    |W| > δ → sign(W), 否则 → 0
    """

    @staticmethod
    def forward(ctx, input, delta):
    ctx.save_for_backward(input, delta)
    out = torch.zeros_like(input)
    out[input > delta] = 1
    out[input < -delta] = -1
    return out

    @staticmethod
    def backward(ctx, grad_output):
    input, delta = ctx.saved_tensors
    grad_input = grad_output.clone()
    # 仅在 [-delta, delta] 外传递梯度
    mask = (input.abs() > delta)
    grad_input[~mask] = 0
    return grad_input, None
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    class TernaryConv2d(nn.Module):
    """三值化卷积层"""

    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0):
    super().__init__()
    self.in_channels = in_channels
    self.out_channels = out_channels
    self.kernel_size = kernel_size
    self.stride = stride
    self.padding = padding

    self.weight = nn.Parameter(torch.randn(out_channels, in_channels, kernel_size, kernel_size))
    self.bias = nn.Parameter(torch.zeros(out_channels))

    def forward(self, x):
    # 计算阈值 δ = 0.7 * mean(|W|)
    delta = 0.7 * self.weight.abs().mean()

    # 三值化权重
    weight_t = TernaryQuantize.apply(self.weight, delta)

    # 前向计算
    out = F.conv2d(x, weight_t, self.bias, self.stride, self.padding)
    return out

执行结果,模型大小: 0.33 MB (原始 FP32 ~3MB)

自动轻量化(AutoML)

使用神经架构搜索(NAS)或强化学习自动寻找最优轻量结构。

  • AMC(AutoML for Model Compression)

  • Once-for-All (OFA):训练一个超网,可从中提取不同大小的子网。

实施步骤

  1. 定义搜索空间(如通道数、层数、算子类型)。

  2. 设计奖励函数(精度 + 延迟约束)。

  3. 使用RL、进化算法或梯度方法搜索。

  4. 训练最优子网并部署。

端到端实施流程

  1. 基准建立

    • 在目标任务上训练一个高精度模型(如ResNet50)。
    • 测量其参数量、FLOPs、推理延迟、内存占用。
  2. 选择轻量化策略组合

    • 优先考虑结构设计 + 量化(通用性强)。
    • 若精度要求高,加入知识蒸馏。
    • 若已有大模型,可尝试剪枝 + QAT。
  3. 迭代优化

    image-20251012153208075
  4. 部署验证

    • 使用目标硬件(如ARM CPU、NPU)测试实际延迟与功耗。
    • 使用TFLite、ONNX、TensorRT等格式转换工具。

GPT2 轻量化实现

通过以下技术对 gpt2 模型的轻量化 demo 实现。

微调恢复精度:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

def fine_tune_pruned(model, tokenizer, num_steps=100):
"""微调剪枝后模型(恢复精度)"""
print(" 开始微调剪枝后模型...")
model.train()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)

# 简单文本用于微调
texts = [
"Artificial intelligence is a wonderful field.",
"Machine learning enables computers to learn from data.",
"Natural language processing allows machines to understand human language.",
"The future of AI is bright and full of possibilities.",
"Deep learning has revolutionized many areas of technology."
] * 20 # 重复以增加数据量

for step in range(num_steps):
text = texts[step % len(texts)]
inputs = tokenizer(text, return_tensors="pt", max_length=128, truncation=True)
input_ids = inputs.input_ids.to(device)

outputs = model(input_ids, labels=input_ids)
loss = outputs.loss

optimizer.zero_grad()
loss.backward()
optimizer.step()

if step % 20 == 0:
print(f" Step {step}, Loss: {loss.item():.4f}")

model.eval()
return model

模型剪枝

剪枝

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
def prune_gpt2_mlp(model: GPT2LMHeadModel, prune_ratio=0.3):
"""
对 GPT-2 的 MLP 层进行结构化剪枝

剪枝策略:
- 评估 c_fc 和 c_proj 的输出通道重要性
- 移除 L1 范数最小的 prune_ratio 比例通道
- 重建模型结构确保维度匹配
"""
print(f" 开始剪枝 GPT-2 MLP 层 (prune_ratio={prune_ratio})...")

for layer_idx, block in enumerate(model.transformer.h):
# print(f" 处理层 {layer_idx}...")

# === 1. 剪枝 c_fc (768 -> 3072) ===
c_fc = block.mlp.c_fc
l1_norm_fc = get_conv1d_l1_norm(c_fc)
num_total_fc = len(l1_norm_fc)
num_keep_fc = int(num_total_fc * (1 - prune_ratio))
_, keep_indices_fc = torch.topk(l1_norm_fc, num_keep_fc, largest=True)
pruned_c_fc = prune_conv1d_layer(c_fc, keep_indices_fc)
block.mlp.c_fc = pruned_c_fc

# === 2. 剪枝 c_proj 的 INPUT 通道(3072 -> num_keep_fc)===
c_proj = block.mlp.c_proj
# c_proj.weight.shape = [3072, 768] → 输入通道是第 0 维
weight_proj = c_proj.weight.data # [in=3072, out=768]
bias_proj = c_proj.bias.data if c_proj.bias is not None else None

# 剪枝输入通道:保留 keep_indices_fc 对应的行
pruned_weight_proj = weight_proj[keep_indices_fc, :] # [num_keep_fc, 768]

# 创建新 c_proj: in=num_keep_fc, out=768
new_c_proj = Conv1D(768, num_keep_fc)
with torch.no_grad():
new_c_proj.weight.copy_(pruned_weight_proj)
if bias_proj is not None:
new_c_proj.bias.copy_(bias_proj) # bias 不变(输出维度仍是 768)

block.mlp.c_proj = new_c_proj

print(" GPT-2 MLP 剪枝完成!")
return model

索引剪枝层:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
def prune_conv1d_layer(conv1d_layer: Conv1D, keep_indices: torch.Tensor):
"""
根据保留的输出通道索引剪枝 Conv1D 层

Args:
conv1d_layer: 原始 Conv1D 层
keep_indices: 要保留的输出通道索引 (shape=[num_keep])

Returns:
新的 Conv1D 层(输出通道数 = num_keep)
"""
weight = conv1d_layer.weight.data # [in, out]
bias = conv1d_layer.bias.data if conv1d_layer.bias is not None else None

# 剪枝权重: 保留 keep_indices 对应的输出通道
pruned_weight = weight[:, keep_indices] # [in, num_keep]

# 创建新层
in_features = weight.shape[0]
out_features = len(keep_indices)
new_layer = Conv1D(out_features, in_features)

with torch.no_grad():
new_layer.weight.copy_(pruned_weight)
if bias is not None:
new_layer.bias.copy_(bias[keep_indices])

return new_layer

量化感知

待完善

知识蒸馏

1
2
3
4
5
# Teacher: gpt2-medium (355M 参数)
teacher = GPT2LMHeadModel.from_pretrained("openai-community/gpt2-medium").to(device)

# Student: gpt2 (124M 参数)
student = GPT2LMHeadModel.from_pretrained("openai-community/gpt2").to(device)

蒸馏损失:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
def compute_language_modeling_loss(logits, labels):
"""标准语言建模损失"""
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
loss = fun.cross_entropy(
shift_logits.view(-1, shift_logits.size(-1)),
shift_labels.view(-1),
ignore_index=-100
)
return loss


def distillation_loss(student_logits, teacher_logits, labels,
alpha=0.5, temperature=1.5):
"""
知识蒸馏损失 = α * 软标签损失 + (1-α) * 硬标签损失

Args:
student_logits: Student 模型输出 [B, T, V]
teacher_logits: Teacher 模型输出 [B, T, V]
labels: 真实标签 [B, T]
alpha: 软标签损失权重(0.5~0.9)
temperature: 温度参数(>1 平滑分布)
"""
# === 1. 软标签损失(KL 散度)===
# Teacher 软化概率分布
teacher_probs = fun.log_softmax(teacher_logits / temperature, dim=-1)
# Student 软化概率分布
student_probs = fun.log_softmax(student_logits / temperature, dim=-1)

# KL 散度损失
kl_loss = fun.kl_div(
student_probs,
teacher_probs,
reduction='batchmean'
) * (temperature ** 2)

# 硬标签损失(高权重)
ce_loss = compute_language_modeling_loss(student_logits, labels)

# 总损失:KL 权重更低
total_loss = alpha * kl_loss + (1 - alpha) * ce_loss

return total_loss, kl_loss, ce_loss

使用以上蒸馏技术后效果:

Student (蒸馏后):The future of AI is,(出现 Token 循环问题)

困惑度对比(500 样本):

  • Teacher (gpt2-medium): 40.44

  • Student (原始 gpt2): 54.64

  • Student (蒸馏后): 123638626614094164131840.00

通过调整参数(alpha、temperature、total_loss)后,效果甚微。

改用三阶段蒸馏后,效果如下:

Student (蒸馏后): The future of AI is still a mystery.

困惑度对比(200 样本):

  • Teacher (gpt2-medium): 35.91

  • Student (原始 gpt2): 49.94

  • Student (蒸馏后): 36.26

低秩分解

低秩分解的核心方法

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
def decompose_conv1d_svd(conv1d_layer: Conv1D, rank_ratio=0.5):
"""
使用奇异值分解(SVD)对 Hugging Face 的 Conv1D 层进行低秩近似分解。

背景说明:
- Hugging Face 的 GPT-2/Transformer 模型使用自定义 Conv1D 层(非标准 nn.Conv1d)
- Conv1D(out_features, in_features) 表示:输入 in_features 维,输出 out_features 维
- 其权重矩阵 W 的形状为 [in_features, out_features](注意:与 nn.Linear 相反!)
- 前向计算公式:output = input @ W + bias (标准矩阵乘法,无需转置)

分解目标:
将原始权重矩阵 W ∈ R^{n×m}(n=in_features, m=out_features)近似分解为:
W ≈ A @ B
其中:
A ∈ R^{n×r} (第一层权重)
B ∈ R^{r×m} (第二层权重)
r = rank_ratio * min(n, m) (低秩近似秩)

优势:
- 原始参数量:n × m
- 分解后参数量:n × r + r × m = r × (n + m)
- 当 r << min(n, m) 时,显著减少参数量和计算量

Args:
conv1d_layer (Conv1D): 待分解的 Hugging Face Conv1D 层
rank_ratio (float): 保留的秩比例(0.0 ~ 1.0),值越小压缩率越高,但精度损失越大

Returns:
nn.Sequential: 由两个 Conv1D 层组成的序列,功能等价于原始层(近似)
"""

# === 步骤 1: 提取原始权重和偏置 ===
# Conv1D.weight 形状: [in_features, out_features]
W = conv1d_layer.weight.data # 获取权重张量(不计算梯度)

# 提取偏置(如果存在)
# Conv1D.bias 形状: [out_features]
bias = conv1d_layer.bias.data if conv1d_layer.bias is not None else None

# === 步骤 2: 获取输入/输出维度 ===
# W.shape = [in_features, out_features]
in_features, out_features = W.shape

# 计算最大可能秩(矩阵的秩不超过 min(行数, 列数))
max_rank = min(in_features, out_features)

# 根据 rank_ratio 计算目标分解秩 r
r = int(rank_ratio * max_rank)
# 确保 r 至少为 1(避免秩为 0)
r = max(1, r)
# 确保 r 不超过 max_rank - 1(避免数值不稳定,且 SVD 需要 r < 秩)
r = min(r, max_rank - 1)

# === 步骤 3: 执行奇异值分解(SVD)===
# 对权重矩阵 W 执行 SVD: W = U @ diag(S) @ Vh
# - U: 左奇异向量矩阵,形状 [in_features, max_rank]
# - S: 奇异值向量,形状 [max_rank]
# - Vh: 右奇异向量矩阵的转置,形状 [max_rank, out_features]
# 使用 full_matrices=False 以节省内存(只计算必要部分)
U, S, Vh = torch.linalg.svd(W, full_matrices=False)

# === 步骤 4: 提取前 r 个奇异分量 ===
# 取前 r 列的左奇异向量: [in_features, r]
U_r = U[:, :r]
# 取前 r 个奇异值: [r]
S_r = S[:r]
# 取前 r 行的右奇异向量转置: [r, out_features]
Vh_r = Vh[:r, :]

# === 步骤 5: 构造低秩近似矩阵 ===
# 数学原理:W ≈ U_r @ diag(S_r) @ Vh_r
# 为数值稳定性和对称性,将奇异值平方根分配到两边:
# A = U_r @ diag(sqrt(S_r)) → [in_features, r]
# B = diag(sqrt(S_r)) @ Vh_r → [r, out_features]
# 这样 A @ B = U_r @ diag(S_r) @ Vh_r ≈ W

# 计算 sqrt(S_r) 并扩展维度以支持广播
sqrt_S = torch.sqrt(S_r) # [r]

# 构造 A = U_r * sqrt(S_r)
# 使用 unsqueeze(0) 将 sqrt_S 变为 [1, r],与 U_r [in_features, r] 广播相乘
A = U_r * sqrt_S.unsqueeze(0) # 结果形状: [in_features, r]

# 构造 B = sqrt(S_r) * Vh_r
# 使用 unsqueeze(1) 将 sqrt_S 变为 [r, 1],与 Vh_r [r, out_features] 广播相乘
B = sqrt_S.unsqueeze(1) * Vh_r # 结果形状: [r, out_features]

# === 步骤 6: 创建两个新的 Conv1D 层 ===
# 第一层: 输入 in_features → 输出 r
# Conv1D(out_features=r, in_features=in_features)
conv1 = Conv1D(r, in_features)

# 第二层: 输入 r → 输出 out_features
# Conv1D(out_features=out_features, in_features=r)
# 注意:偏置只加在最后一层(与原始层一致)
conv2 = Conv1D(out_features, r)

# === 步骤 7: 复制分解后的权重和偏置 ===
with torch.no_grad(): # 禁用梯度计算,避免影响优化器状态
# 复制第一层权重 A ([in_features, r])
# conv1.weight 形状应为 [in_features, r]
conv1.weight.copy_(A)

# 复制第二层权重 B ([r, out_features])
# conv2.weight 形状应为 [r, out_features]
conv2.weight.copy_(B)

# 复制偏置(仅第二层需要,因为原始偏置作用于最终输出)
if bias is not None:
# conv2.bias 形状: [out_features]
conv2.bias.copy_(bias)

# === 步骤 8: 返回组合层 ===
# 使用 nn.Sequential 将两个 Conv1D 层串联
# 前向计算: input -> conv1 -> conv2 -> output
# 功能等价于: output = input @ W + bias (近似)
return nn.Sequential(conv1, conv2)

分解 + 生成对比

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
def main():
model_name = "openai-community/gpt2"

# 1. 加载原始 GPT-2
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
model_orig = GPT2LMHeadModel.from_pretrained(model_name).to(device)
model_orig.eval()

print(f"原始模型参数量: {count_parameters(model_orig):.2f} M")

# 2. 复制模型用于分解
model_decomp = GPT2LMHeadModel.from_pretrained(model_name).to(device)
model_decomp.eval()

# 3. 低秩分解 0.3 重复; 0.5 可用; 0.6 较好; 0.7 接近原始
# 分解全部 + 微调
# model_decomp = decompose_gpt2_mlp(model_decomp, rank_ratio=0.5)
# model_decomp = fine_tune_decomposed(model_decomp, tokenizer, num_steps=60)

# 推荐:只分解 c_proj + rank_ratio=0.6
model_decomp = decompose_gpt2_mlp(model_decomp, rank_ratio=0.6, decompose_c_fc=False)

print(f"分解后模型参数量: {count_parameters(model_decomp):.2f} M")

# 4. 文本生成对比
prompt = "The future of AI is"
print(f"\n 提示词: '{prompt}'")

text_orig = generate_text(model_orig, tokenizer, prompt)
print(f"\n 原始模型:\n{text_orig}")

text_decomp = generate_text(model_decomp, tokenizer, prompt)
print(f"\n 分解模型:\n{text_decomp}")

# 5. 保存分解后模型
model_decomp.save_pretrained("../../data/gpt2_decomposed")
tokenizer.save_pretrained("../../data/gpt2_decomposed")

不同场景的轻量化

同一模型在不同场景下的轻量化策略应有所不同。轻量化不是“一刀切”的技术,而是需要根据部署硬件、应用场景、性能约束和业务需求进行定制化设计。

轻量化的本质是在 精度(Accuracy)速度(Latency)体积(Size)功耗(Power) 之间做权衡。不同场景对这些指标的优先级完全不同。

场景 硬件平台 关键约束 推荐轻量化技术 原因
1. 手机端实时语音助手 ARM CPU / NPU 低延迟(<300ms)
中等精度
1.量化(INT8)
2.结构化剪枝
3.轻量架构(Squeezeformer)
- CPU/NPU 对 INT8 有硬件加速
- 延迟敏感,需移除冗余计算
- 精度可小幅牺牲(WER +0.5% 可接受)
2. 智能音箱(离线唤醒) MCU / DSP(<100MHz) 极低功耗
模型 <1MB
低精度
1.二值化/三值化
2.知识蒸馏(Tiny Student)
3.算子融合 + 固定点
- 内存极小,需极致压缩
- 唤醒词任务简单,小模型足够
- 浮点运算耗电,需定点
3. 车载语音系统 Automotive SoC(如 Qualcomm SA8155) 高可靠性
实时性(<500ms)
中高精度
1.量化感知训练(QAT)
2.通道剪枝 + 微调
3.模型分割(CPU+NPU)
- 安全关键,精度损失需 <0.3%
- SoC 有专用 NPU,需 INT8 优化
- 需支持多语言,模型不能太小
4. 服务器高并发 API x86 CPU / GPU 高吞吐(QPS)
低延迟
1.TensorRT INT8(GPU)
2.ONNX + 并行推理
3.动态批处理
- GPU 用 TensorRT 效果远超 PTQ
- CPU 用 OpenVINO / ONNX Runtime
- 吞吐优先,可接受稍大模型
5. IoT 传感器(关键词检测) Cortex-M 系列 <100KB 模型
<10mW 功耗
1.MCU 专用框架(TensorFlow)
2.手工设计 Tiny CNN
3.无浮点,全整型
- 无操作系统,需静态内存
- 模型必须 <100KB
- 通常只检测 1~10 个关键词

场景 1:手机 App(Android)

  • 目标:实时转录,延迟 <500ms

  • 策略:

    • 使用 WeNet + Squeezeformer
    • INT8 量化(TFLite + NNAPI)
    • 剪枝 30% 通道
    • 模型大小:~5MB
    • WER:+0.8%

场景 2:智能手表(Wear OS)

  • 目标:离线命令识别,模型 <2MB

  • 策略:

    • 知识蒸馏:Teacher=Conformer, Student=Tiny-Conformer(4 层)
    • FP16 量化(无 INT8 支持)
    • 移除语言模型(仅 CTC)
    • 模型大小:~1.8MB
    • WER:+2.5%(可接受,因命令简单)

场景 3:车载系统(Linux + NPU)

  • 目标:多语言支持,高鲁棒性

  • 策略:

    • QAT(量化感知训练)
    • 仅剪枝 10%(保留精度)
    • TensorRT 导出
    • 模型大小:~8MB
    • WER:+0.2%
技术 适用场景 不适用场景
量化(PTQ/QAT) 手机、服务器、车载(有 INT8 支持) MCU(无 SIMD 指令)
结构化剪枝 通用(CPU/GPU/NPU) 超小模型(剪枝收益低)
知识蒸馏 需要小模型 + 有大 Teacher 无预训练大模型
低秩分解 大 Linear 层(如 FFN) 小卷积层(如 3x3)
轻量架构设计 从零训练 已有大模型需压缩
二值化/三值化 MCU、超低功耗 高精度任务(如医疗 ASR)

常用工具与框架

技术 工具/库
剪枝 NNI, TorchPruner, TensorFlow Model Optimization
量化 PyTorch Quantization, TensorFlow Lite, TensorRT
蒸馏 HuggingFace Transformers (DistilBERT), TorchDistill
轻量网络 timm (PyTorch Image Models), torchvision
自动压缩 AutoKeras, OFA official repo