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


环境声明

  • 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. 理解模型压缩的必要性和挑战
  2. 掌握模型量化技术(INT8、FP16、GPTQ、AWQ等)
  3. 深入理解模型剪枝的原理和方法
  4. 应用知识蒸馏训练小模型
  5. 了解2024-2026年最新压缩技术进展
  6. 使用TensorRT和ONNX Runtime部署优化模型
  7. 在实际项目中应用模型压缩技术

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(qZ)

其中:

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

流程

  1. 训练原始模型至收敛
  2. 剪枝一定比例的权重
  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 量化注意事项

  1. 校准数据选择

    • 使用与推理数据分布相似的校准数据
    • 校准数据量不宜过少(通常100-1000个样本)
  2. 量化感知训练

    • 使用较小的学习率
    • 训练轮数可以适当减少
  3. 动态范围

    • 注意处理异常值(outliers)
    • 可以使用per-channel量化提高精度

8.2 剪枝注意事项

  1. 剪枝比例

    • 不要一次性剪枝过多
    • 建议每次剪枝10-30%
  2. 重训练

    • 剪枝后必须重训练
    • 使用较小的学习率微调
  3. 结构化剪枝

    • 注意层间的维度匹配
    • 可能需要重新构建模型结构

8.3 部署注意事项

  1. 硬件兼容性

    • 确认目标硬件支持的操作
    • 不同硬件可能有不同的优化策略
  2. 批处理大小

    • 选择合适的batch size
    • 大batch size通常能获得更好的吞吐
  3. 预热

    • 首次推理通常较慢
    • 建议进行预热后再测试性能

9. 本章小结

核心知识点回顾

模型量化

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

模型剪枝

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

知识蒸馏

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

部署优化

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

压缩技术选择指南

场景 推荐技术 理由
快速部署 PTQ 无需重新训练
精度优先 QAT 精度损失最小
大语言模型 GPTQ/AWQ 专用优化
边缘设备 结构化剪枝+量化 易于硬件加速
极致压缩 蒸馏+剪枝+量化 组合使用

10. 思考与练习

基础练习

  1. 量化计算:手动计算一个权重矩阵的INT8量化过程,包括缩放因子和零点的计算。

  2. 剪枝分析:分析不同剪枝粒度(非结构化、结构化)的优缺点,说明在什么场景下应该选择哪种方法。

  3. 蒸馏温度:解释知识蒸馏中温度参数的作用,为什么需要调整温度。

进阶练习

  1. 混合精度训练:实现一个支持FP16训练但保持FP32权重的混合精度训练流程。

  2. 动态量化推理:为Transformer模型实现动态量化推理,并比较与静态量化的性能差异。

  3. 通道剪枝:实现基于BN层gamma因子的通道剪枝方法。

挑战练习

  1. 大模型量化:使用AutoGPTQ或AutoAWQ对一个7B参数的语言模型进行4-bit量化,并评估推理速度和精度。

  2. 端到端优化:对一个ResNet模型进行完整的压缩流程(剪枝+量化+蒸馏),目标是在精度损失<1%的情况下实现4倍压缩。

  3. 移动端部署:将一个图像分类模型部署到Android设备上,使用TensorFlow Lite或PyTorch Mobile。


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


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

Logo

腾讯云面向开发者汇聚海量精品云计算使用和开发经验,营造开放的云计算技术生态圈。

更多推荐