深度学习模型压缩与加速：从理论到工程实践

1次阅读

没有评论

共计 2623 个字符，预计需要花费 7 分钟才能阅读完成。

近年来，深度学习模型的规模呈指数级增长。以自然语言处理为例，GPT- 3 的参数量达到了惊人的 1750 亿。这种趋势带来了两个主要问题：

计算资源消耗：大模型训练和推理需要大量 GPU/TPU 资源，成本极高
推理延迟：在移动设备或边缘计算场景中，大模型难以满足实时性要求

根据我们的实测数据，ResNet-50 在 CPU 上的推理延迟达到 120ms/ 帧，难以满足视频流实时处理的需求。这就是模型压缩技术存在的意义。

目前主流的模型压缩方法可分为三类：

原理：移除模型中不重要的权重或通道
优点：可大幅减少参数量和计算量
缺点：需要精细调参避免精度损失
适用场景：CNN 类模型，特别是视觉任务

原理：将 FP32 权重转换为低精度(如 INT8) 表示
优点：内存占用减少 4 倍，计算速度提升 2 - 4 倍
缺点：极端量化(如 4 -bit) 可能导致精度下降
适用场景：所有需要部署的模型

原理：用大模型(教师) 指导小模型 (学生) 训练
优点：可保持较高模型精度
缺点：需要额外训练过程
适用场景：需要保持高精度的场景

以下是一个完整的通道剪枝实现：

import torch
import torch.nn as nn
import torch.nn.utils.prune as prune

class CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, 3)
        self.conv2 = nn.Conv2d(64, 128, 3)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        return x

# 1. 初始化模型
model = CNN()

# 2. 选择剪枝方法（这里使用 L1 范数剪枝）parameters_to_prune = ((model.conv1, 'weight'),
    (model.conv2, 'weight'),
)

# 3. 执行剪枝（剪掉 50% 的通道）prune.global_unstructured(
    parameters_to_prune,
    pruning_method=prune.L1Unstructured,
    amount=0.5,
)

# 4. 永久移除剪枝的权重
for module, param in parameters_to_prune:
    prune.remove(module, param)

关键点说明：

使用 PyTorch 内置的 prune 模块实现标准化剪枝
L1Unstructured 按权重绝对值大小进行剪枝
amount=0.5 表示剪枝 50% 的通道

PyTorch 提供了两种量化方式：

import torch.quantization

# 原始模型
model_fp32 = CNN()

# 转换为量化模型
model_int8 = torch.quantization.quantize_dynamic(
    model_fp32,  # 原始模型
    {nn.Linear, nn.Conv2d},  # 要量化的模块类型
    dtype=torch.qint8  # 量化数据类型
)

# 需要准备校准数据集
calibration_data = torch.randn(100, 3, 224, 224)

# 量化配置
model_fp32.eval()
model_fp32.qconfig = torch.quantization.get_default_qconfig('fbgemm')

# 准备量化模型
model_prepared = torch.quantization.prepare(model_fp32)

# 校准（确定量化参数）model_prepared(calibration_data)

# 转换为最终量化模型
model_int8 = torch.quantization.convert(model_prepared)

典型的知识蒸馏流程包含三个关键组件：

# 1. 教师模型（大模型）teacher = BigModel()
teacher.load_state_dict(torch.load('teacher.pth'))

# 2. 学生模型（小模型）student = SmallModel()

# 3. 蒸馏损失函数
criterion = nn.KLDivLoss()
optimizer = torch.optim.Adam(student.parameters())

# 训练循环
for inputs, labels in dataloader:
    # 教师预测
    with torch.no_grad():
        teacher_logits = teacher(inputs)

    # 学生预测
    student_logits = student(inputs)

    # 计算损失
    loss = criterion(F.log_softmax(student_logits/T, dim=1),
        F.softmax(teacher_logits/T, dim=1)
    ) * (T*T) + F.cross_entropy(student_logits, labels)

    # 反向传播
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

其中 T 是温度参数，通常设为 2 -5。

我们在不同硬件上测试了 ResNet-50 的加速效果：