【机器学习精通】第20章 | 模型压缩与加速:让大模型跑得更快
摘要:本文系统介绍了深度学习模型压缩与加速的关键技术,包括量化、剪枝和知识蒸馏等方法。量化部分详细讲解了训练后量化(PTQ)和量化感知训练(QAT)的实现,并涵盖了大语言模型专用量化技术(GPTQ、AWQ)。文章提供了PyTorch代码示例,展示了从模型准备到实际量化的完整流程,帮助开发者在资源受限环境下实现高效模型部署。压缩技术可显著减少模型大小(2-10倍)并提升推理速度,同时保持可接受的精度
模型压缩与加速是部署深度学习模型的关键技术。本章将系统讲解量化、剪枝、知识蒸馏等核心方法,让你的模型在资源受限环境下高效运行。
环境声明
- Python版本:Python 3.12+
- 核心依赖:PyTorch 2.0+、TensorRT、ONNX、transformers
- 开发工具:PyCharm / VS Code / Jupyter Notebook
- 操作系统:Windows / macOS / Linux(通用)
# 安装依赖
pip install torch torchvision onnx onnxruntime transformers
pip install neural-compressor # Intel模型压缩工具
学习目标
完成本章学习后,你将能够:
- 理解模型压缩的必要性和挑战
- 掌握模型量化技术(INT8、FP16、GPTQ、AWQ等)
- 深入理解模型剪枝的原理和方法
- 应用知识蒸馏训练小模型
- 了解2024-2026年最新压缩技术进展
- 使用TensorRT和ONNX Runtime部署优化模型
- 在实际项目中应用模型压缩技术
1. 模型压缩概述
1.1 为什么需要模型压缩
部署挑战:
| 场景 | 限制 | 典型需求 |
|---|---|---|
| 移动端APP | 内存<4GB,存储<100MB | 实时推理 |
| 边缘设备 | 算力有限,功耗敏感 | 低延迟 |
| 云端服务 | 成本优化 | 高吞吐 |
| 自动驾驶 | 安全关键,实时性 | 确定性延迟 |
压缩收益:
- 减少模型存储空间(通常减少4-8倍)
- 降低内存占用
- 加速推理速度(通常加速2-4倍)
- 降低能耗
1.2 模型压缩技术分类
| 技术类别 | 原理 | 压缩比 | 精度损失 |
|---|---|---|---|
| 量化(Quantization) | 降低权重精度 | 2-4x | 低 |
| 剪枝(Pruning) | 移除不重要权重 | 2-10x | 中 |
| 知识蒸馏(Distillation) | 大模型教小模型 | 2-10x | 中 |
| 低秩分解(Low-rank) | 矩阵分解 | 2-4x | 低 |
| 神经架构搜索(NAS) | 搜索高效架构 | 2-5x | 低 |
1.3 压缩-精度权衡
帕累托最优:
在模型大小和精度之间寻找最佳平衡点。通常:
- 量化:首选方法,精度损失小
- 剪枝:配合重训练使用
- 蒸馏:需要教师模型
2. 模型量化
2.1 量化基础
什么是量化
将浮点数权重和激活值转换为低精度表示(如INT8)。
量化公式:
r=S(q−Z) r = S(q - Z) r=S(q−Z)
其中:
- rrr:原始浮点值
- qqq:量化后的整数值
- SSS:缩放因子(scale)
- ZZZ:零点(zero point)
对称量化 vs 非对称量化:
| 类型 | 公式 | 适用场景 |
|---|---|---|
| 对称量化 | q=round(r/S)q = \text{round}(r / S)q=round(r/S) | 权重(以0为中心) |
| 非对称量化 | q=round(r/S)+Zq = \text{round}(r / S) + Zq=round(r/S)+Z | 激活值(非对称分布) |
2.2 训练后量化(PTQ)
核心思想
在模型训练完成后进行量化,无需重新训练。
动态量化:
import torch
import torch.quantization
# 动态量化(仅量化权重,激活值动态计算)
model_fp32 = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
model_int8 = torch.quantization.quantize_dynamic(
model_fp32, # 原始模型
{torch.nn.Linear}, # 要量化的层类型
dtype=torch.qint8 # 量化类型
)
# 比较模型大小
import os
torch.save(model_fp32.state_dict(), "model_fp32.pth")
torch.save(model_int8.state_dict(), "model_int8.pth")
fp32_size = os.path.getsize("model_fp32.pth") / 1024 / 1024
int8_size = os.path.getsize("model_int8.pth") / 1024 / 1024
print(f"FP32模型大小: {fp32_size:.2f} MB")
print(f"INT8模型大小: {int8_size:.2f} MB")
print(f"压缩比: {fp32_size / int8_size:.2f}x")
静态量化:
import torch
from torch.quantization import quantize_fx
# 准备模型
model_fp32 = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
model_fp32.eval()
# 准备校准数据
def calibrate(model, data_loader):
model.eval()
with torch.no_grad():
for images, _ in data_loader:
model(images)
# 静态量化
model_prepared = quantize_fx.prepare_fx(model_fp32, {'': 'fbgemm'})
calibrate(model_prepared, calib_loader) # 使用校准数据
model_int8 = quantize_fx.convert_fx(model_prepared)
2.3 量化感知训练(QAT)
核心思想
在训练过程中模拟量化效果,让模型适应量化带来的误差。
import torch
import torch.nn as nn
from torch.quantization import QuantStub, DeQuantStub
class QuantizableModel(nn.Module):
"""可量化模型示例"""
def __init__(self):
super().__init__()
self.quant = QuantStub() # 量化输入
self.conv1 = nn.Conv2d(3, 64, 3)
self.relu = nn.ReLU()
self.fc = nn.Linear(64, 10)
self.dequant = DeQuantStub() # 反量化输出
def forward(self, x):
x = self.quant(x)
x = self.conv1(x)
x = self.relu(x)
x = x.mean([2, 3]) # 全局平均池化
x = self.fc(x)
x = self.dequant(x)
return x
# 量化感知训练流程
model = QuantizableModel()
model.train()
# 插入伪量化节点
model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
model_prepared = torch.quantization.prepare_qat(model)
# 训练(伪量化节点会模拟量化效果)
optimizer = torch.optim.Adam(model_prepared.parameters())
for epoch in range(num_epochs):
for images, labels in train_loader:
optimizer.zero_grad()
outputs = model_prepared(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# 转换为量化模型
model_int8 = torch.quantization.convert(model_prepared.eval())
2.4 大语言模型量化(GPTQ、AWQ、GGUF)
GPTQ(2022-2023)
逐层量化,使用OBS(Optimal Brain Surgeon)方法最小化量化误差。
# 使用AutoGPTQ进行量化
from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
model_name = "meta-llama/Llama-2-7b-hf"
quantize_config = BaseQuantizeConfig(
bits=4, # 4-bit量化
group_size=128,
desc_act=False,
)
model = AutoGPTQForCausalLM.from_pretrained(
model_name,
quantize_config
)
# 准备校准数据
calib_data = ["样本文本1", "样本文本2", ...]
# 量化
model.quantize(calib_data)
# 保存
model.save_quantized("llama-7b-4bit-gptq")
AWQ(2023-2024)
激活感知权重量化,根据激活值的重要性进行保护。
# 使用AutoAWQ进行量化
from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer
model_name = "meta-llama/Llama-2-7b-hf"
quant_config = {
"zero_point": True,
"q_group_size": 128,
"w_bit": 4,
"version": "GEMM"
}
# 加载模型
model = AutoAWQForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
# 准备校准数据
examples = [
tokenizer("auto-gptq is an easy-to-use model quantization library", return_tensors="pt")
]
# 量化
model.quantize(
tokenizer,
quant_config=quant_config,
calib_data=examples
)
# 保存
model.save_quantized("llama-7b-4bit-awq")
GGUF/GGML(llama.cpp)
针对CPU推理优化的量化格式。
# 使用llama.cpp转换和量化
# 1. 转换为GGUF格式
python convert-hf-to-gguf.py models/llama-7b
# 2. 量化(多种量化方案)
./quantize models/llama-7b-f16.gguf models/llama-7b-q4_0.gguf q4_0
./quantize models/llama-7b-f16.gguf models/llama-7b-q5_k_m.gguf q5_k_m
| 量化类型 | 每权重位数 | 适用场景 |
|---|---|---|
| Q4_0 | 4.5 | 最小模型,精度损失较大 |
| Q5_0 | 5.5 | 平衡选择 |
| Q8_0 | 8.5 | 精度优先 |
| Q4_K_M | 4.7 | 推荐方案 |
| Q5_K_M | 5.7 | 高精度需求 |
3. 模型剪枝
3.1 剪枝基础
什么是剪枝
移除神经网络中不重要的权重或神经元,减少模型复杂度。
剪枝粒度:
| 粒度 | 描述 | 压缩比 | 实现难度 |
|---|---|---|---|
| 非结构化剪枝 | 移除单个权重 | 高 | 难(需要稀疏库支持) |
| 结构化剪枝 | 移除整个通道/层 | 中 | 易 |
| 半结构化剪枝 | 移除块稀疏模式 | 中 | 中 |
3.2 非结构化剪枝
基于幅度的剪枝
移除绝对值最小的权重。
import torch
import torch.nn.utils.prune as prune
# 定义模型
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
# 对卷积层进行剪枝(剪枝30%的权重)
for name, module in model.named_modules():
if isinstance(module, torch.nn.Conv2d):
prune.l1_unstructured(module, name='weight', amount=0.3)
# 使剪枝永久化
for name, module in model.named_modules():
if isinstance(module, torch.nn.Conv2d):
prune.remove(module, 'weight')
# 检查稀疏度
def check_sparsity(model):
total_params = 0
zero_params = 0
for param in model.parameters():
total_params += param.numel()
zero_params += (param == 0).sum().item()
return zero_params / total_params
sparsity = check_sparsity(model)
print(f"模型稀疏度: {sparsity:.2%}")
3.3 结构化剪枝
通道剪枝
移除不重要的整个通道。
import torch
import torch.nn as nn
def prune_channels(conv_layer, prune_ratio=0.3):
"""
基于L1范数的通道剪枝
参数:
conv_layer: 卷积层
prune_ratio: 剪枝比例
"""
weight = conv_layer.weight.data
# 计算每个通道的重要性(L1范数)
channel_importance = weight.abs().sum(dim=(1, 2, 3))
# 确定要保留的通道数
num_channels = weight.size(0)
num_keep = int(num_channels * (1 - prune_ratio))
# 选择最重要的通道
_, keep_indices = torch.topk(channel_importance, num_keep)
# 创建新的卷积层
new_conv = nn.Conv2d(
in_channels=conv_layer.in_channels,
out_channels=num_keep,
kernel_size=conv_layer.kernel_size,
stride=conv_layer.stride,
padding=conv_layer.padding,
bias=conv_layer.bias is not None
)
# 复制权重
new_conv.weight.data = weight[keep_indices]
if conv_layer.bias is not None:
new_conv.bias.data = conv_layer.bias.data[keep_indices]
return new_conv, keep_indices
# 应用结构化剪枝
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
# 对第一个卷积层进行剪枝
new_conv, keep_indices = prune_channels(model.conv1, prune_ratio=0.3)
model.conv1 = new_conv
print(f"原始通道数: 64, 剪枝后通道数: {new_conv.out_channels}")
3.4 迭代剪枝与重训练
流程:
- 训练原始模型至收敛
- 剪枝一定比例的权重
- 重训练恢复精度
- 重复步骤2-3直到达到目标稀疏度
def iterative_pruning(model, train_loader, val_loader,
target_sparsity=0.8, prune_steps=5, epochs_per_step=10):
"""
迭代剪枝流程
"""
sparsity_per_step = target_sparsity / prune_steps
current_sparsity = 0
for step in range(prune_steps):
print(f"\n剪枝步骤 {step + 1}/{prune_steps}")
# 1. 剪枝
current_sparsity += sparsity_per_step
for name, module in model.named_modules():
if isinstance(module, torch.nn.Conv2d):
prune.l1_unstructured(module, name='weight', amount=current_sparsity)
print(f"当前稀疏度: {current_sparsity:.2%}")
# 2. 重训练
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
for epoch in range(epochs_per_step):
model.train()
for images, labels in train_loader:
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# 3. 评估
model.eval()
correct = 0
total = 0
with torch.no_grad():
for images, labels in val_loader:
outputs = model(images)
_, predicted = outputs.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
accuracy = 100. * correct / total
print(f"重训练后精度: {accuracy:.2f}%")
return model
4. 知识蒸馏
4.1 知识蒸馏基础
核心思想
让大模型(教师模型)指导小模型(学生模型)学习,使学生模型获得接近教师模型的性能。
为什么有效:
- 软标签包含更多类别间关系信息
- 教师模型的"暗知识"(如错误类别的相对概率)对学生有帮助
4.2 基本蒸馏方法
软标签蒸馏
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.7):
super().__init__()
self.temperature = temperature
self.alpha = alpha # 蒸馏损失的权重
def forward(self, student_logits, teacher_logits, labels):
"""
student_logits: 学生模型的输出
teacher_logits: 教师模型的输出
labels: 真实标签
"""
# 硬标签损失(交叉熵)
hard_loss = F.cross_entropy(student_logits, labels)
# 软标签损失(KL散度)
soft_student = F.log_softmax(student_logits / self.temperature, dim=1)
soft_teacher = F.softmax(teacher_logits / self.temperature, dim=1)
soft_loss = F.kl_div(soft_student, soft_teacher, reduction='batchmean')
soft_loss = soft_loss * (self.temperature ** 2)
# 总损失
loss = self.alpha * soft_loss + (1 - self.alpha) * hard_loss
return loss
# 蒸馏训练流程
def train_with_distillation(teacher, student, train_loader, val_loader,
num_epochs=50, temperature=4.0, alpha=0.7):
"""
使用知识蒸馏训练学生模型
"""
teacher.eval() # 教师模型保持评估模式
student.train()
criterion = DistillationLoss(temperature=temperature, alpha=alpha)
optimizer = torch.optim.Adam(student.parameters(), lr=1e-3)
for epoch in range(num_epochs):
student.train()
total_loss = 0
for images, labels in train_loader:
optimizer.zero_grad()
# 教师模型预测(无梯度)
with torch.no_grad():
teacher_logits = teacher(images)
# 学生模型预测
student_logits = student(images)
# 计算蒸馏损失
loss = criterion(student_logits, teacher_logits, labels)
loss.backward()
optimizer.step()
total_loss += loss.item()
# 验证
student.eval()
correct = 0
total = 0
with torch.no_grad():
for images, labels in val_loader:
outputs = student(images)
_, predicted = outputs.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
accuracy = 100. * correct / total
avg_loss = total_loss / len(train_loader)
print(f"Epoch [{epoch+1}/{num_epochs}] "
f"Loss: {avg_loss:.4f} | "
f"Accuracy: {accuracy:.2f}%")
return student
4.3 特征蒸馏
核心思想
不仅蒸馏输出,还蒸馏中间层特征。
class FeatureDistillationLoss(nn.Module):
"""特征蒸馏损失"""
def __init__(self, temperature=4.0, alpha=0.7, beta=0.3):
super().__init__()
self.temperature = temperature
self.alpha = alpha # 软标签权重
self.beta = beta # 特征蒸馏权重
def forward(self, student_logits, teacher_logits, labels,
student_features, teacher_features):
"""
student_features: 学生模型的中间特征
teacher_features: 教师模型的中间特征
"""
# 硬标签损失
hard_loss = F.cross_entropy(student_logits, labels)
# 软标签损失
soft_student = F.log_softmax(student_logits / self.temperature, dim=1)
soft_teacher = F.softmax(teacher_logits / self.temperature, dim=1)
soft_loss = F.kl_div(soft_student, soft_teacher, reduction='batchmean')
soft_loss *= self.temperature ** 2
# 特征蒸馏损失(MSE)
# 可能需要先对齐特征维度
feature_loss = F.mse_loss(student_features, teacher_features)
# 总损失
loss = (1 - self.alpha - self.beta) * hard_loss + \
self.alpha * soft_loss + \
self.beta * feature_loss
return loss
4.4 自蒸馏与在线蒸馏
自蒸馏(Self-Distillation)
模型自己教自己,不需要单独的教师模型。
在线蒸馏(Online Distillation)
多个学生模型相互学习,如Deep Mutual Learning。
5. 模型部署优化
5.1 ONNX导出与优化
导出模型:
import torch
import torch.onnx
# 加载模型
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
model.eval()
# 创建虚拟输入
dummy_input = torch.randn(1, 3, 224, 224)
# 导出ONNX
torch.onnx.export(
model,
dummy_input,
"resnet18.onnx",
export_params=True,
opset_version=11,
do_constant_folding=True,
input_names=['input'],
output_names=['output'],
dynamic_axes={
'input': {0: 'batch_size'},
'output': {0: 'batch_size'}
}
)
print("模型已导出为ONNX格式")
ONNX Runtime推理:
import onnxruntime as ort
import numpy as np
# 创建推理会话
session = ort.InferenceSession("resnet18.onnx")
# 获取输入输出信息
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
# 准备输入数据
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
# 推理
outputs = session.run([output_name], {input_name: input_data})
print(f"输出形状: {outputs[0].shape}")
5.2 TensorRT优化
TensorRT是NVIDIA的高性能推理库。
import torch
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
class TensorRTInference:
"""TensorRT推理引擎"""
def __init__(self, onnx_path, fp16_mode=True):
self.logger = trt.Logger(trt.Logger.WARNING)
# 构建TensorRT引擎
builder = trt.Builder(self.logger)
network = builder.create_network(
1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
)
parser = trt.OnnxParser(network, self.logger)
# 解析ONNX模型
with open(onnx_path, 'rb') as f:
parser.parse(f.read())
# 配置构建器
config = builder.create_builder_config()
config.max_workspace_size = 1 << 30 # 1GB
if fp16_mode:
config.set_flag(trt.BuilderFlag.FP16)
# 构建引擎
self.engine = builder.build_engine(network, config)
self.context = self.engine.create_execution_context()
# 分配GPU内存
self.allocate_buffers()
def allocate_buffers(self):
"""分配输入输出缓冲区"""
self.inputs = []
self.outputs = []
self.bindings = []
for binding in self.engine:
shape = self.engine.get_binding_shape(binding)
size = trt.volume(shape)
dtype = trt.nptype(self.engine.get_binding_dtype(binding))
# 分配GPU内存
host_mem = cuda.pagelocked_empty(size, dtype)
device_mem = cuda.mem_alloc(host_mem.nbytes)
self.bindings.append(int(device_mem))
if self.engine.binding_is_input(binding):
self.inputs.append({'host': host_mem, 'device': device_mem})
else:
self.outputs.append({'host': host_mem, 'device': device_mem})
def infer(self, input_data):
"""执行推理"""
# 复制输入数据到GPU
np.copyto(self.inputs[0]['host'], input_data.ravel())
cuda.memcpy_htod(self.inputs[0]['device'], self.inputs[0]['host'])
# 执行推理
self.context.execute_v2(bindings=self.bindings)
# 复制输出数据到CPU
for out in self.outputs:
cuda.memcpy_dtoh(out['host'], out['device'])
return [out['host'] for out in self.outputs]
# 使用示例
trt_engine = TensorRTInference("resnet18.onnx", fp16_mode=True)
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
outputs = trt_engine.infer(input_data)
5.3 TorchScript与TorchServe
TorchScript:
import torch
# 加载模型
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
model.eval()
# 追踪模型(Tracing)
example_input = torch.randn(1, 3, 224, 224)
traced_model = torch.jit.trace(model, example_input)
# 保存
torch.jit.save(traced_model, "resnet18_traced.pt")
# 加载并推理
loaded_model = torch.jit.load("resnet18_traced.pt")
output = loaded_model(example_input)
6. 2024-2026年最新压缩技术
6.1 稀疏注意力(Sparse Attention)
核心思想
降低Transformer的注意力计算复杂度。
| 方法 | 复杂度 | 代表工作 |
|---|---|---|
| 局部注意力 | O(n) | Sparse Transformer |
| 线性注意力 | O(n) | Performer、Linformer |
| 稀疏模式 | O(n log n) | Longformer、BigBird |
| 滑动窗口 | O(n) | Swin Transformer |
6.2 混合专家模型(MoE)
核心思想
只激活部分参数,降低推理成本。
class MoELayer(nn.Module):
"""混合专家层"""
def __init__(self, d_model, num_experts=8, top_k=2):
super().__init__()
self.num_experts = num_experts
self.top_k = top_k
# 专家网络
self.experts = nn.ModuleList([
nn.Sequential(
nn.Linear(d_model, d_model * 4),
nn.GELU(),
nn.Linear(d_model * 4, d_model)
) for _ in range(num_experts)
])
# 门控网络
self.gate = nn.Linear(d_model, num_experts)
def forward(self, x):
"""
x: [batch, seq_len, d_model]
"""
batch_size, seq_len, d_model = x.shape
x_flat = x.view(-1, d_model) # [batch*seq_len, d_model]
# 计算门控权重
gate_logits = self.gate(x_flat) # [batch*seq_len, num_experts]
weights, indices = torch.topk(F.softmax(gate_logits, dim=-1), self.top_k)
weights = weights / weights.sum(dim=-1, keepdim=True)
# 只激活top-k专家
output = torch.zeros_like(x_flat)
for i in range(self.top_k):
expert_idx = indices[:, i]
expert_weight = weights[:, i:i+1]
for j in range(self.num_experts):
mask = expert_idx == j
if mask.any():
expert_input = x_flat[mask]
expert_output = self.experts[j](expert_input)
output[mask] += expert_weight[mask] * expert_output
return output.view(batch_size, seq_len, d_model)
6.3 推测解码(Speculative Decoding)
核心思想
使用小模型快速生成候选token,大模型验证并修正。
加速效果:
- 通常可获得2-3倍的加速
- 不改变输出分布
6.4 模型压缩工具链(2024-2026)
| 工具 | 功能 | 支持框架 |
|---|---|---|
| AutoGPTQ | GPTQ量化 | PyTorch |
| AutoAWQ | AWQ量化 | PyTorch |
| llama.cpp | GGUF量化与推理 | C++ |
| vLLM | PagedAttention推理 | PyTorch |
| TensorRT-LLM | NVIDIA GPU优化 | TensorRT |
| ONNX Runtime | 跨平台推理 | ONNX |
| Intel Neural Compressor | Intel硬件优化 | PyTorch/TensorFlow |
7. 实战案例:压缩BERT模型
7.1 完整压缩流程
import torch
from transformers import BertModel, BertTokenizer
import torch.nn as nn
import torch.nn.utils.prune as prune
# 加载原始模型
model_name = 'bert-base-uncased'
tokenizer = BertTokenizer.from_pretrained(model_name)
teacher_model = BertModel.from_pretrained(model_name)
print(f"原始模型参数量: {sum(p.numel() for p in teacher_model.parameters()) / 1e6:.2f}M")
# 1. 量化
quantized_model = torch.quantization.quantize_dynamic(
teacher_model,
{torch.nn.Linear},
dtype=torch.qint8
)
# 2. 剪枝(非结构化)
for name, module in quantized_model.named_modules():
if isinstance(module, torch.nn.Linear):
prune.l1_unstructured(module, name='weight', amount=0.3)
# 3. 知识蒸馏训练小模型
class TinyBERT(nn.Module):
"""小型BERT模型"""
def __init__(self, vocab_size=30522, hidden_size=312, num_layers=4):
super().__init__()
self.embeddings = teacher_model.embeddings
# 减少层数和隐藏维度
encoder_layer = nn.TransformerEncoderLayer(
d_model=hidden_size,
nhead=12,
dim_feedforward=hidden_size * 4,
batch_first=True
)
self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
# 投影层匹配教师维度
self.projection = nn.Linear(768, hidden_size)
self.inverse_projection = nn.Linear(hidden_size, 768)
def forward(self, input_ids, attention_mask=None):
embedding_output = self.embeddings(input_ids)
embedding_output = self.projection(embedding_output)
if attention_mask is not None:
attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
attention_mask = (1.0 - attention_mask) * -10000.0
encoder_output = self.encoder(embedding_output, src_key_padding_mask=attention_mask)
encoder_output = self.inverse_projection(encoder_output)
return encoder_output
# 初始化学生模型
student_model = TinyBERT()
# 蒸馏训练
optimizer = torch.optim.Adam(student_model.parameters(), lr=2e-5)
criterion = nn.MSELoss()
# 模拟训练
for epoch in range(3):
student_model.train()
# 模拟输入
inputs = tokenizer("This is a sample text", return_tensors="pt")
# 教师模型输出
with torch.no_grad():
teacher_output = teacher_model(**inputs).last_hidden_state
# 学生模型输出
student_output = student_model(inputs['input_ids'], inputs['attention_mask'])
# 蒸馏损失
loss = criterion(student_output, teacher_output)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")
# 评估压缩效果
print(f"\n压缩后模型参数量: {sum(p.numel() for p in student_model.parameters()) / 1e6:.2f}M")
print(f"压缩比: {sum(p.numel() for p in teacher_model.parameters()) / sum(p.numel() for p in student_model.parameters()):.2f}x")
8. 避坑小贴士
8.1 量化注意事项
-
校准数据选择
- 使用与推理数据分布相似的校准数据
- 校准数据量不宜过少(通常100-1000个样本)
-
量化感知训练
- 使用较小的学习率
- 训练轮数可以适当减少
-
动态范围
- 注意处理异常值(outliers)
- 可以使用per-channel量化提高精度
8.2 剪枝注意事项
-
剪枝比例
- 不要一次性剪枝过多
- 建议每次剪枝10-30%
-
重训练
- 剪枝后必须重训练
- 使用较小的学习率微调
-
结构化剪枝
- 注意层间的维度匹配
- 可能需要重新构建模型结构
8.3 部署注意事项
-
硬件兼容性
- 确认目标硬件支持的操作
- 不同硬件可能有不同的优化策略
-
批处理大小
- 选择合适的batch size
- 大batch size通常能获得更好的吞吐
-
预热
- 首次推理通常较慢
- 建议进行预热后再测试性能
9. 本章小结
核心知识点回顾
模型量化:
- 训练后量化(PTQ):快速但精度损失可能较大
- 量化感知训练(QAT):精度更高但需要重新训练
- 大模型量化:GPTQ、AWQ、GGUF等专用方法
模型剪枝:
- 非结构化剪枝:压缩比高但加速困难
- 结构化剪枝:易于实现加速
- 迭代剪枝:逐步达到目标稀疏度
知识蒸馏:
- 软标签蒸馏:利用教师模型的概率分布
- 特征蒸馏:蒸馏中间层特征
- 自蒸馏:无需单独的教师模型
部署优化:
- ONNX:跨平台标准格式
- TensorRT:NVIDIA GPU优化
- TorchScript:PyTorch生产部署
压缩技术选择指南
| 场景 | 推荐技术 | 理由 |
|---|---|---|
| 快速部署 | PTQ | 无需重新训练 |
| 精度优先 | QAT | 精度损失最小 |
| 大语言模型 | GPTQ/AWQ | 专用优化 |
| 边缘设备 | 结构化剪枝+量化 | 易于硬件加速 |
| 极致压缩 | 蒸馏+剪枝+量化 | 组合使用 |
10. 思考与练习
基础练习
-
量化计算:手动计算一个权重矩阵的INT8量化过程,包括缩放因子和零点的计算。
-
剪枝分析:分析不同剪枝粒度(非结构化、结构化)的优缺点,说明在什么场景下应该选择哪种方法。
-
蒸馏温度:解释知识蒸馏中温度参数的作用,为什么需要调整温度。
进阶练习
-
混合精度训练:实现一个支持FP16训练但保持FP32权重的混合精度训练流程。
-
动态量化推理:为Transformer模型实现动态量化推理,并比较与静态量化的性能差异。
-
通道剪枝:实现基于BN层gamma因子的通道剪枝方法。
挑战练习
-
大模型量化:使用AutoGPTQ或AutoAWQ对一个7B参数的语言模型进行4-bit量化,并评估推理速度和精度。
-
端到端优化:对一个ResNet模型进行完整的压缩流程(剪枝+量化+蒸馏),目标是在精度损失<1%的情况下实现4倍压缩。
-
移动端部署:将一个图像分类模型部署到Android设备上,使用TensorFlow Lite或PyTorch Mobile。
学习建议:模型压缩是工程实践性很强的领域。建议读者在实际项目中尝试不同的压缩方法,理解各种技术的适用场景和限制。
本系列教程持续更新中,欢迎关注专栏获取更多机器学习深度内容。
更多推荐
所有评论(0)