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


环境声明

  • 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_04.5最小模型,精度损失较大
Q5_05.5平衡选择
Q8_08.5精度优先
Q4_K_M4.7推荐方案
Q5_K_M5.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)

工具功能支持框架
AutoGPTQGPTQ量化PyTorch
AutoAWQAWQ量化PyTorch
llama.cppGGUF量化与推理C++
vLLMPagedAttention推理PyTorch
TensorRT-LLMNVIDIA GPU优化TensorRT
ONNX Runtime跨平台推理ONNX
Intel Neural CompressorIntel硬件优化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

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

更多推荐