[模型优化] 3. 模型压缩

👋 你好!这里有实用干货与深度分享✨✨ 若有帮助,欢迎:​
👍 点赞 | ⭐ 收藏 | 💬 评论 | ➕ 关注 ,解锁更多精彩!​
📁 收藏专栏即可第一时间获取最新推送🔔。​
📖后续我将持续带来更多优质内容,期待与你一同探索知识,携手前行,共同进步🚀。​



人工智能

模型压缩

本文详细介绍深度学习模型压缩的相关技术,包括权重剪枝、知识蒸馏、量化、低秩分解和结构化稀疏化等方法,帮助您在保持模型性能的同时减小模型体积、降低计算复杂度。

1. 权重剪枝

1.1 剪枝策略

剪枝类型原理特点适用场景
幅值剪枝基于权重绝对值大小,设定阈值进行筛选保留重要连接适合稠密网络
敏感度剪枝基于权重对损失的影响,计算权重重要性剪枝不重要的连接,计算开销较大对精度要求高的场景
结构化剪枝按通道、滤波器整体剪枝保持规则结构,硬件友好,实际加速效果好需要实际部署加速的场景

1.2 实现示例

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

# 创建示例模型
model = torch.nn.Linear(10, 10)

# 1. 基于L1范数的非结构化剪枝
prune.l1_unstructured(model, name='weight', amount=0.3)  # 剪掉30%的权重

# 2. 结构化剪枝(按维度)
prune.ln_structured(model, name='weight', amount=0.5, n=2, dim=0)  # 剪掉50%的输出神经元

# 3. 自定义剪枝方法
class CustomPruning(prune.BasePruningMethod):
    def compute_mask(self, t, default_mask):
        mask = default_mask.clone()
        mask.view(-1)[torch.abs(t.view(-1)) < 0.1] = 0  # 剪掉绝对值小于0.1的权重
        return mask

# 应用自定义剪枝
CustomPruning.apply(model, 'weight')

# 4. 移除剪枝(使权重永久性变为0)
prune.remove(model, 'weight')

# 5. 迭代式剪枝示例
def iterative_pruning(model, total_prune_ratio, steps):
    ratio_per_step = 1 - (1 - total_prune_ratio) ** (1 / steps)
    for i in range(steps):
        # 对每一层进行剪枝
        for name, module in model.named_modules():
            if isinstance(module, torch.nn.Linear) or isinstance(module, torch.nn.Conv2d):
                prune.l1_unstructured(module, name='weight', amount=ratio_per_step)
        
        # 训练几个epoch恢复精度
        train_model(model, epochs=5)
    
    # 最终移除剪枝
    for name, module in model.named_modules():
        if isinstance(module, torch.nn.Linear) or isinstance(module, torch.nn.Conv2d):
            prune.remove(module, 'weight')

2. 知识蒸馏

2.1 蒸馏方法

蒸馏类型原理特点适用场景
响应蒸馏学习教师模型的输出分布使用软标签进行训练
平衡硬标签和软标签
分类任务
特征蒸馏学习中间层特征对齐特征表示
传递更多知识
复杂模型
关系蒸馏学习样本间的关系保持数据结构
增强泛化能力
需要保持数据关系的场景

2.2 实现示例

import torch
import torch.nn.functional as F

class DistillationLoss(torch.nn.Module):
    def __init__(self, alpha=0.5, temperature=2.0):
        super().__init__()
        self.alpha = alpha  # 硬标签和软标签损失的平衡系数
        self.T = temperature  # 温度参数,控制软标签的平滑程度
        
    def forward(self, student_outputs, teacher_outputs, targets):
        # 硬标签损失
        hard_loss = F.cross_entropy(student_outputs, targets)
        
        # 软标签损失
        soft_loss = F.kl_div(
            F.log_softmax(student_outputs / self.T, dim=1),
            F.softmax(teacher_outputs / self.T, dim=1),
            reduction='batchmean'
        ) * (self.T * self.T)  # 乘以T²进行缩放
        
        # 总损失
        return self.alpha * hard_loss + (1 - self.alpha) * soft_loss

# 特征蒸馏损失
class FeatureDistillationLoss(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.mse_loss = torch.nn.MSELoss()
    
    def forward(self, student_features, teacher_features):
        # 可能需要调整特征维度
        if student_features.shape != teacher_features.shape:
            # 使用自适应池化调整空间维度
            student_features = F.adaptive_avg_pool2d(
                student_features, 
                output_size=teacher_features.shape[2:]
            )
        return self.mse_loss(student_features, teacher_features)

# 训练循环
distillation_criterion = DistillationLoss(alpha=0.5, temperature=4.0)
feature_criterion = FeatureDistillationLoss()

for epoch in range(num_epochs):
    for data, targets in train_loader:
        # 教师模型推理
        with torch.no_grad():
            teacher_outputs, teacher_features = teacher_model(data, return_features=True)
        
        # 学生模型训练
        student_outputs, student_features = student_model(data, return_features=True)
        
        # 计算损失
        response_loss = distillation_criterion(student_outputs, teacher_outputs, targets)
        feature_loss = feature_criterion(student_features, teacher_features)
        
        # 总损失
        loss = response_loss + 0.1 * feature_loss
        loss.backward()
        optimizer.step()

3. 量化

3.1 量化方法

量化方法原理特点适用场景
训练后量化直接将FP32权重转换为INT8无需重新训练
精度损失较大
对精度要求不高的场景
量化感知训练训练过程中模拟量化效果权重适应量化误差
精度损失较小
需要保持较高精度的场景
混合精度量化不同层使用不同精度敏感层保持高精度
不敏感层使用低精度
平衡性能和精度
对某些层精度要求高的场景

3.2 实现示例

import torch

# 1. 训练后静态量化
def post_training_quantization(model):
    # 准备量化配置
    model.qconfig = torch.quantization.get_default_qconfig('fbgemm')  # x86架构
    # model.qconfig = torch.quantization.get_default_qconfig('qnnpack')  # ARM架构
    
    # 准备模型进行量化
    model_prepared = torch.quantization.prepare(model)
    
    # 校准(使用一小部分数据)
    with torch.no_grad():
        for data, _ in calibration_loader:  # 使用校准数据集
            model_prepared(data)
    
    # 转换为量化模型
    quantized_model = torch.quantization.convert(model_prepared)
    
    return quantized_model

# 2. 量化感知训练
class QAT_Model(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model
        # 设置量化配置
        self.model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
        # 准备QAT
        self.model_prepared = torch.quantization.prepare_qat(self.model)
    
    def forward(self, x):
        return self.model_prepared(x)
    
    def convert(self):
        # 训练结束后转换为量化模型
        return torch.quantization.convert(self.model_prepared.eval())

# 使用示例
model = create_model()  # 创建浮点模型
qat_model = QAT_Model(model)

# 训练循环
for epoch in range(num_epochs):
    for data, targets in train_loader:
        outputs = qat_model(data)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

# 转换为最终量化模型
quantized_model = qat_model.convert()

4. 低秩分解

4.1 分解方法

分解方法原理特点适用场景
SVD分解奇异值分解保留主要成分
重构权重矩阵
适合全连接层
CP分解张量分解多维数据压缩
降低参数量
适合高维卷积核
Tucker分解高阶SVD保留多维结构适合多通道卷积

4.2 实现示例

import torch
import numpy as np
import tensorly as tl
from tensorly.decomposition import parafac, tucker

# 设置后端
tl.set_backend('pytorch')

# 1. SVD分解(适用于全连接层)
def decompose_linear_layer(layer, rank):
    # 获取权重
    weight = layer.weight.data
    
    # SVD分解
    U, S, Vh = torch.linalg.svd(weight, full_matrices=False)
    
    # 截断到指定秩
    U = U[:, :rank]
    S = S[:rank]
    Vh = Vh[:rank, :]
    
    # 创建新层
    first_layer = torch.nn.Linear(weight.shape[1], rank, bias=False)
    second_layer = torch.nn.Linear(rank, weight.shape[0], bias=True)
    
    # 设置权重
    first_layer.weight.data = Vh
    second_layer.weight.data = U * S.view(-1, 1)
    if layer.bias is not None:
        second_layer.bias.data = layer.bias.data
    
    return torch.nn.Sequential(first_layer, second_layer)

# 2. CP分解(适用于卷积层)
def decompose_conv_layer_cp(layer, rank):
    # 获取权重
    weight = layer.weight.data
    
    # 重塑为张量
    weight_tensor = weight.reshape(weight.shape[0], weight.shape[1], -1)
    
    # CP分解
    factors = parafac(weight_tensor, rank=rank, init='random')
    
    # 创建分解后的层
    pointwise_s_to_r = torch.nn.Conv2d(
        in_channels=weight.shape[1],
        out_channels=rank,
        kernel_size=1,
        stride=1,
        padding=0,
        bias=False
    )
    
    depthwise_r = torch.nn.Conv2d(
        in_channels=rank,
        out_channels=rank,
        kernel_size=layer.kernel_size,
        stride=layer.stride,
        padding=layer.padding,
        groups=rank,
        bias=False
    )
    
    pointwise_r_to_t = torch.nn.Conv2d(
        in_channels=rank,
        out_channels=weight.shape[0],
        kernel_size=1,
        stride=1,
        padding=0,
        bias=layer.bias is not None
    )
    
    # 设置权重(简化实现)
    # 实际应用中需要根据CP分解结果正确设置权重
    
    if layer.bias is not None:
        pointwise_r_to_t.bias.data = layer.bias.data
    
    return torch.nn.Sequential(
        pointwise_s_to_r,
        depthwise_r,
        pointwise_r_to_t
    )

# 使用示例
original_linear = torch.nn.Linear(1000, 1000)
decomposed_linear = decompose_linear_layer(original_linear, rank=100)

original_conv = torch.nn.Conv2d(64, 64, kernel_size=3, padding=1)
decomposed_conv = decompose_conv_layer_cp(original_conv, rank=32)

5. 结构化稀疏化

5.1 稀疏化方法

稀疏化方法原理特点适用场景
组稀疏按组进行稀疏化保持硬件友好结构
提高实际加速效果
适合并行计算
通道剪枝整通道级别剪枝降低模型复杂度
减少计算量
无需特殊硬件支持
需要直接加速的场景
结构化规范化使用结构化正则项引导权重形成规则模式
便于硬件加速
需要硬件加速的场景

5.2 实现示例

import torch

# 1. 组稀疏正则化
class GroupSparseLoss(torch.nn.Module):
    def __init__(self, lambda_l1=0.01):
        super().__init__()
        self.lambda_l1 = lambda_l1
        
    def forward(self, model):
        group_loss = 0
        for name, param in model.named_parameters():
            if 'weight' in name:
                # 计算组L2,1正则化(按行L2范数的L1和)
                group_loss += torch.sum(torch.sqrt(torch.sum(param.pow(2), dim=1) + 1e-8))
        return self.lambda_l1 * group_loss

# 2. 通道重要性评估与剪枝
def prune_channels(model, prune_ratio=0.3):
    # 收集所有卷积层
    conv_layers = []
    for name, module in model.named_modules():
        if isinstance(module, torch.nn.Conv2d):
            conv_layers.append((name, module))
    
    for name, layer in conv_layers:
        # 计算每个通道的L1范数
        channel_norms = torch.sum(torch.abs(layer.weight), dim=[1, 2, 3])
        
        # 确定阈值
        num_channels = layer.weight.shape[0]
        num_to_prune = int(num_channels * prune_ratio)
        threshold = torch.sort(channel_norms)[0][num_to_prune]
        
        # 创建掩码
        mask = (channel_norms > threshold).float().view(-1, 1, 1, 1)
        
        # 应用掩码
        layer.weight.data = layer.weight.data * mask
        
        # 如果有偏置,也应用掩码
        if layer.bias is not None:
            layer.bias.data = layer.bias.data * mask.view(-1)

# 3. 结构化训练示例
criterion = torch.nn.CrossEntropyLoss()
group_sparse = GroupSparseLoss(lambda_l1=0.001)

for epoch in range(num_epochs):
    for data, target in train_loader:
        output = model(data)
        # 任务损失
        task_loss = criterion(output, target)
        # 稀疏化损失
        sparse_loss = group_sparse(model)
        # 总损失
        loss = task_loss + sparse_loss
        loss.backward()
        optimizer.step()
    
    # 每隔几个epoch进行一次剪枝
    if epoch % 5 == 0 and epoch > 0:
        prune_channels(model, prune_ratio=0.1)
        # 微调几个epoch
        finetune(model, epochs=2)

6. 压缩方法对比

压缩方法压缩比精度损失推理加速训练成本适用场景
非结构化剪枝需专用硬件过参数化模型
结构化剪枝直接加速冗余通道多的CNN
知识蒸馏直接加速有大模型作教师
量化(INT8)4倍直接加速低/中推理部署场景
低秩分解直接加速全连接层、卷积层
结构化稀疏直接加速需规则结构的场景

7. 最佳实践

7.1 压缩策略选择

  • 模型分析

    • 分析模型结构特点
    • 识别参数冗余部分
    • 确定瓶颈层
    • 针对性选择压缩方法
  • 部署环境考量

    • 考虑硬件限制(内存、计算能力)
    • 评估延迟要求
    • 选择硬件友好的压缩方法
    • 测试不同方法的实际加速效果
  • 多方法组合

    • 剪枝+量化组合使用
    • 知识蒸馏辅助其他压缩方法
    • 不同层使用不同压缩策略
    • 逐步应用多种方法

7.2 压缩过程控制

  • 渐进式压缩

    • 小步迭代压缩
    • 每次压缩后微调
    • 监控性能变化
    • 避免一次性大幅压缩
  • 性能监控

    • 设置性能指标阈值
    • 监控关键指标变化
    • 建立压缩率-精度曲线
    • 找到最佳平衡点
  • 自动化搜索

    • 使用AutoML方法
    • 自动搜索最佳压缩策略
    • 优化压缩超参数
    • 减少人工调参工作

7.3 精度恢复

  • 微调策略

    • 压缩后使用小学习率微调
    • 考虑使用学习率预热
    • 适当延长训练时间
    • 使用原始预训练权重初始化
  • 知识蒸馏辅助

    • 使用原始模型作为教师
    • 指导压缩模型学习
    • 保留原始模型的泛化能力
    • 缓解精度下降
  • 数据增强

    • 增加训练数据多样性
    • 使用更强的数据增强
    • 考虑使用生成式方法扩充数据
    • 提高模型鲁棒性

7.4 实际部署

  • 部署前验证

    • 在目标硬件上测试
    • 验证实际加速效果
    • 确认内存占用
    • 检查精度是否符合预期
  • 硬件适配

    • 针对特定硬件优化
    • 利用硬件加速特性
    • 考虑量化精度兼容性
    • 测试不同批量大小
  • 持续优化

    • 收集实际使用数据
    • 分析性能瓶颈
    • 针对性进一步优化
    • 考虑模型更新策略

8. 常见问题与解决方案

  1. 精度下降严重

    • 减小压缩率
    • 增加微调轮数
    • 使用知识蒸馏辅助训练
    • 尝试更精细的压缩粒度
  2. 实际加速效果不明显

    • 检查是否使用结构化压缩
    • 确认硬件是否支持相应优化
    • 分析推理瓶颈是否在被压缩部分
    • 考虑端到端优化而非单一层优化
  3. 量化模型不稳定

    • 使用量化感知训练
    • 对敏感层使用更高精度
    • 增加校准数据集大小和多样性
    • 检查是否有极值激活
  4. 剪枝后模型结构混乱

    • 使用结构化剪枝
    • 保持层间依赖关系
    • 使用通道级别剪枝
    • 考虑整体网络结构
  5. 知识蒸馏效果不佳

    • 调整温度参数
    • 平衡软标签和硬标签权重
    • 增加中间层特征蒸馏
    • 使用更强大的教师模型

通过合理选择和组合模型压缩技术,可以在保持模型性能的同时,显著减小模型体积、降低计算复杂度,使深度学习模型更适合在资源受限的环境中部署。



📌 感谢阅读!若文章对你有用,别吝啬互动~​
👍 点个赞 | ⭐ 收藏备用 | 💬 留下你的想法 ,关注我,更多干货持续更新!

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

斌zz

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值