【机器学习精通】第20章 | 模型压缩与加速：让大模型跑得更快

模型压缩与加速是部署深度学习模型的关键技术。本章将系统讲解量化、剪枝、知识蒸馏等核心方法，让你的模型在资源受限环境下高效运行。

环境声明

• 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模型压缩工具

学习目标

完成本章学习后，你将能够：

1. 模型压缩概述

1.1 为什么需要模型压缩

部署挑战：

压缩收益：

• 减少模型存储空间（通常减少4-8倍）
• 降低内存占用
• 加速推理速度（通常加速2-4倍）
• 降低能耗

1.2 模型压缩技术分类

1.3 压缩-精度权衡

帕累托最优：

在模型大小和精度之间寻找最佳平衡点。通常：

• 量化：首选方法，精度损失小
• 剪枝：配合重训练使用
• 蒸馏：需要教师模型

2. 模型量化

2.1 量化基础

什么是量化

将浮点数权重和激活值转换为低精度表示（如INT8）。

量化公式：

r

=

S

(

q

−

Z

)

r = S(q – Z)

r=S(q−Z)

其中：

• $r$
• $q$
• $S$
• $Z$

对称量化 vs 非对称量化：

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

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 迭代剪枝与重训练

流程：

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的注意力计算复杂度。

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）

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 量化注意事项

8.2 剪枝注意事项

8.3 部署注意事项

9. 本章小结

核心知识点回顾

模型量化：

• 训练后量化（PTQ）：快速但精度损失可能较大
• 量化感知训练（QAT）：精度更高但需要重新训练
• 大模型量化：GPTQ、AWQ、GGUF等专用方法

模型剪枝：

• 非结构化剪枝：压缩比高但加速困难
• 结构化剪枝：易于实现加速
• 迭代剪枝：逐步达到目标稀疏度

知识蒸馏：

• 软标签蒸馏：利用教师模型的概率分布
• 特征蒸馏：蒸馏中间层特征
• 自蒸馏：无需单独的教师模型

部署优化：

• ONNX：跨平台标准格式
• TensorRT：NVIDIA GPU优化
• TorchScript：PyTorch生产部署

压缩技术选择指南

10. 思考与练习

基础练习

进阶练习

挑战练习

学习建议：模型压缩是工程实践性很强的领域。建议读者在实际项目中尝试不同的压缩方法，理解各种技术的适用场景和限制。

本系列教程持续更新中，欢迎关注专栏获取更多机器学习深度内容。