当前位置：首页 > news >正文

CNN 模型压缩：剪枝、量化与知识蒸馏

news 2026/4/15 20:55:11

CNN 模型压缩：剪枝、量化与知识蒸馏

核心结论

剪枝：移除冗余权重，减少模型参数量和计算量
量化：降低权重和激活值的精度，减少存储和计算开销
知识蒸馏：将大型模型的知识迁移到小型模型
性能对比：不同压缩方法在精度、速度和模型大小方面各有优势
组合策略：多种压缩方法结合使用效果更佳

一、模型压缩的必要性

1.1 深度学习模型的挑战

计算资源需求：大型 CNN 模型需要大量计算资源
存储开销：模型文件大小限制了部署场景
推理速度：实时应用对推理速度有严格要求
能耗：移动设备和边缘设备的能耗限制

1.2 模型压缩的目标

减少参数量：降低模型存储需求
减少计算量：提高推理速度
保持精度：在压缩的同时不显著降低模型性能
适配部署环境：使模型能够在资源受限设备上运行

二、剪枝技术

2.1 剪枝的基本原理

非结构化剪枝：随机移除单个权重
结构化剪枝：移除整个神经元或通道
基于重要性：根据权重的重要性决定是否剪枝
迭代剪枝：多次剪枝和微调的过程

2.2 剪枝方法分类

幅度剪枝：基于权重绝对值大小
梯度剪枝：基于权重梯度信息
运动剪枝：基于权重更新的幅度
L1 正则化：通过正则化促进权重稀疏

2.3 代码示例：基于幅度的剪枝

import torch import torch.nn as nn import torch.nn.functional as F # 定义一个简单的 CNN 模型 class SimpleCNN(nn.Module): def __init__(self): super(SimpleCNN, self).__init__() self.conv1 = nn.Conv2d(3, 16, 3, padding=1) self.conv2 = nn.Conv2d(16, 32, 3, padding=1) self.fc1 = nn.Linear(32 * 8 * 8, 128) self.fc2 = nn.Linear(128, 10) def forward(self, x): x = F.relu(self.conv1(x)) x = F.max_pool2d(x, 2) x = F.relu(self.conv2(x)) x = F.max_pool2d(x, 2) x = x.view(-1, 32 * 8 * 8) x = F.relu(self.fc1(x)) x = self.fc2(x) return x # 初始化模型 model = SimpleCNN() # 模拟训练后的模型 for param in model.parameters(): torch.nn.init.normal_(param, mean=0, std=0.1) # 计算模型参数量 def count_parameters(model): return sum(p.numel() for p in model.parameters() if p.requires_grad) print(f"原始模型参数量: {count_parameters(model):,}") # 基于幅度的剪枝函数 def prune_model(model, pruning_ratio): # 收集所有权重 weights = [] for name, param in model.named_parameters(): if 'weight' in name: weights.append((name, param)) # 计算所有权重的绝对值 all_weights = torch.cat([w.view(-1) for _, w in weights]) # 计算阈值 threshold = torch.quantile(torch.abs(all_weights), pruning_ratio) # 执行剪枝 for name, param in weights: mask = torch.abs(param) > threshold param.data *= mask.float() return threshold # 执行剪枝 pruning_ratio = 0.5 # 剪枝 50% 的权重 threshold = prune_model(model, pruning_ratio) # 计算剪枝后的非零参数量 def count_nonzero_parameters(model): count = 0 for param in model.parameters(): if param.requires_grad: count += torch.count_nonzero(param).item() return count print(f"剪枝后非零参数量: {count_nonzero_parameters(model):,}") print(f"剪枝阈值: {threshold:.4f}") print(f"剪枝比例: {(1 - count_nonzero_parameters(model)/count_parameters(model)):.2f}")

2.4 剪枝的挑战与解决方案

精度下降：通过微调恢复精度
硬件加速：结构化剪枝更有利于硬件加速
剪枝粒度：不同粒度的剪枝效果不同
自动化剪枝：使用 NAS 等方法自动寻找最佳剪枝策略

三、量化技术

3.1 量化的基本原理

动态量化：仅量化权重，激活值在运行时量化
静态量化：同时量化权重和激活值，需要校准
感知量化：在训练过程中考虑量化误差
量化感知训练：通过训练减少量化误差

3.2 量化位宽

INT8 量化：最常用的量化方法，精度损失较小
INT4 量化：更高压缩率，但精度损失较大
二进制量化：极端压缩，仅用 1 位表示权重
三值量化：在二进制量化基础上增加零值

3.3 代码示例：PyTorch 量化

import torch import torch.nn as nn import torch.quantization as quant # 定义一个简单的 CNN 模型 class SimpleCNN(nn.Module): def __init__(self): super(SimpleCNN, self).__init__() self.quant = quant.QuantStub() self.conv1 = nn.Conv2d(3, 16, 3, padding=1) self.relu1 = nn.ReLU() self.pool1 = nn.MaxPool2d(2) self.conv2 = nn.Conv2d(16, 32, 3, padding=1) self.relu2 = nn.ReLU() self.pool2 = nn.MaxPool2d(2) self.fc1 = nn.Linear(32 * 8 * 8, 128) self.relu3 = nn.ReLU() self.fc2 = nn.Linear(128, 10) self.dequant = quant.DeQuantStub() def forward(self, x): x = self.quant(x) x = self.conv1(x) x = self.relu1(x) x = self.pool1(x) x = self.conv2(x) x = self.relu2(x) x = self.pool2(x) x = x.view(-1, 32 * 8 * 8) x = self.fc1(x) x = self.relu3(x) x = self.fc2(x) x = self.dequant(x) return x # 初始化模型 model = SimpleCNN() # 准备校准数据 calibration_data = torch.randn(100, 3, 32, 32) # 动态量化 model_dynamic = quant.quantize_dynamic( model, {nn.Linear, nn.Conv2d}, dtype=torch.qint8 ) # 静态量化 model_static = SimpleCNN() model_static.eval() model_static.qconfig = quant.get_default_qconfig('fbgemm') quant.prepare(model_static, inplace=True) # 校准 with torch.no_grad(): for i in range(10): batch = calibration_data[i*10:(i+1)*10] model_static(batch) # 转换 model_static = quant.convert(model_static, inplace=True) # 计算模型大小 def get_model_size(model): import os import tempfile with tempfile.NamedTemporaryFile(suffix='.pth', delete=False) as f: torch.save(model.state_dict(), f) size = os.path.getsize(f.name) os.unlink(f.name) return size print(f"原始模型大小: {get_model_size(model)/1024/1024:.2f} MB") print(f"动态量化模型大小: {get_model_size(model_dynamic)/1024/1024:.2f} MB") print(f"静态量化模型大小: {get_model_size(model_static)/1024/1024:.2f} MB") # 测试推理速度 import time def measure_inference_time(model, input_data, iterations=100): model.eval() start_time = time.time() with torch.no_grad(): for _ in range(iterations): model(input_data) end_time = time.time() return (end_time - start_time) / iterations input_data = torch.randn(1, 3, 32, 32) original_time = measure_inference_time(model, input_data) dynamic_time = measure_inference_time(model_dynamic, input_data) static_time = measure_inference_time(model_static, input_data) print(f"原始模型推理时间: {original_time*1000:.2f} ms") print(f"动态量化模型推理时间: {dynamic_time*1000:.2f} ms") print(f"静态量化模型推理时间: {static_time*1000:.2f} ms")

3.4 量化的挑战与解决方案

精度损失：使用量化感知训练减少损失
硬件支持：不同硬件对量化的支持程度不同
动态范围：处理激活值的动态范围变化
混合精度：对不同层使用不同的量化精度

四、知识蒸馏

4.1 知识蒸馏的基本原理

教师-学生框架：大型教师模型指导小型学生模型
软标签：教师模型的概率分布包含更多信息
温度参数：控制软标签的平滑程度
蒸馏损失：结合软标签损失和硬标签损失

4.2 知识蒸馏方法

传统蒸馏：使用教师模型的软标签训练学生模型
特征蒸馏：使用教师模型的中间特征训练学生模型
关系蒸馏：蒸馏样本之间的关系信息
自蒸馏：模型自己蒸馏自己，无需单独的教师模型

4.3 代码示例：知识蒸馏实现

import torch import torch.nn as nn import torch.optim as optim import torch.nn.functional as F # 定义教师模型（较大的模型） class TeacherModel(nn.Module): def __init__(self): super(TeacherModel, self).__init__() self.conv1 = nn.Conv2d(3, 32, 3, padding=1) self.conv2 = nn.Conv2d(32, 64, 3, padding=1) self.fc1 = nn.Linear(64 * 8 * 8, 256) self.fc2 = nn.Linear(256, 128) self.fc3 = nn.Linear(128, 10) def forward(self, x): x = F.relu(self.conv1(x)) x = F.max_pool2d(x, 2) x = F.relu(self.conv2(x)) x = F.max_pool2d(x, 2) x = x.view(-1, 64 * 8 * 8) x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = self.fc3(x) return x # 定义学生模型（较小的模型） class StudentModel(nn.Module): def __init__(self): super(StudentModel, self).__init__() self.conv1 = nn.Conv2d(3, 16, 3, padding=1) self.conv2 = nn.Conv2d(16, 32, 3, padding=1) self.fc1 = nn.Linear(32 * 8 * 8, 128) self.fc2 = nn.Linear(128, 10) def forward(self, x): x = F.relu(self.conv1(x)) x = F.max_pool2d(x, 2) x = F.relu(self.conv2(x)) x = F.max_pool2d(x, 2) x = x.view(-1, 32 * 8 * 8) x = F.relu(self.fc1(x)) x = self.fc2(x) return x # 初始化模型 teacher = TeacherModel() student = StudentModel() # 模拟教师模型已经训练完成 for param in teacher.parameters(): torch.nn.init.normal_(param, mean=0, std=0.1) # 知识蒸馏训练 class DistillationLoss(nn.Module): def __init__(self, temperature=2.0, alpha=0.5): super(DistillationLoss, self).__init__() self.temperature = temperature self.alpha = alpha self.cross_entropy = nn.CrossEntropyLoss() def forward(self, student_outputs, teacher_outputs, labels): # 计算硬标签损失 hard_loss = self.cross_entropy(student_outputs, labels) # 计算软标签损失 soft_teacher = F.softmax(teacher_outputs / self.temperature, dim=1) soft_student = F.log_softmax(student_outputs / self.temperature, dim=1) soft_loss = F.kl_div(soft_student, soft_teacher, reduction='batchmean') * (self.temperature ** 2) # 组合损失 return self.alpha * hard_loss + (1 - self.alpha) * soft_loss # 准备数据 train_data = torch.randn(1000, 3, 32, 32) train_labels = torch.randint(0, 10, (1000,)) # 优化器和损失函数 optimizer = optim.Adam(student.parameters(), lr=0.001) distillation_loss = DistillationLoss(temperature=2.0, alpha=0.5) # 训练学生模型 student.train() teacher.eval() for epoch in range(10): running_loss = 0.0 for i in range(0, 1000, 32): batch_data = train_data[i:i+32] batch_labels = train_labels[i:i+32] optimizer.zero_grad() # 教师模型输出 with torch.no_grad(): teacher_outputs = teacher(batch_data) # 学生模型输出 student_outputs = student(batch_data) # 计算损失 loss = distillation_loss(student_outputs, teacher_outputs, batch_labels) # 反向传播 loss.backward() optimizer.step() running_loss += loss.item() print(f"Epoch {epoch+1}, Loss: {running_loss/31:.4f}") # 计算模型参数量 print(f"教师模型参数量: {sum(p.numel() for p in teacher.parameters() if p.requires_grad):,}") print(f"学生模型参数量: {sum(p.numel() for p in student.parameters() if p.requires_grad):,}")

4.4 知识蒸馏的挑战与解决方案

教师模型选择：选择合适的教师模型
温度参数调整：找到最佳温度参数
损失函数设计：平衡硬标签和软标签损失
计算开销：训练过程需要同时运行教师和学生模型

五、性能对比实验

5.1 不同压缩方法的性能对比

压缩方法	模型大小	推理速度	精度损失	适用场景
原始模型	100%	100%	0%	资源充足场景
剪枝 (50%)	~50%	~120%	<1%	通用压缩场景
量化 (INT8)	~25%	~150%	<1%	边缘设备部署
知识蒸馏	~30%	~130%	<2%	需要保持精度的场景
组合方法	~10%	~200%	<3%	资源受限场景

5.2 实际模型压缩案例

import torch import torchvision.models as models import torch.quantization as quant # 加载预训练的 ResNet18 模型 model = models.resnet18(pretrained=True) model.eval() # 计算原始模型大小 def get_model_size(model): import os import tempfile with tempfile.NamedTemporaryFile(suffix='.pth', delete=False) as f: torch.save(model.state_dict(), f) size = os.path.getsize(f.name) os.unlink(f.name) return size print(f"原始 ResNet18 模型大小: {get_model_size(model)/1024/1024:.2f} MB") # 动态量化 model_dynamic = quant.quantize_dynamic( model, {torch.nn.Linear, torch.nn.Conv2d}, dtype=torch.qint8 ) print(f"动态量化模型大小: {get_model_size(model_dynamic)/1024/1024:.2f} MB") # 测试推理速度 import time def measure_inference_time(model, input_data, iterations=100): model.eval() start_time = time.time() with torch.no_grad(): for _ in range(iterations): model(input_data) end_time = time.time() return (end_time - start_time) / iterations input_data = torch.randn(1, 3, 224, 224) original_time = measure_inference_time(model, input_data) dynamic_time = measure_inference_time(model_dynamic, input_data) print(f"原始模型推理时间: {original_time*1000:.2f} ms") print(f"动态量化模型推理时间: {dynamic_time*1000:.2f} ms") print(f"速度提升: {original_time/dynamic_time:.2f}x")

5.3 压缩方法组合效果

# 组合剪枝和量化 import torch import torch.nn as nn import torch.quantization as quant # 定义模型 class SimpleCNN(nn.Module): def __init__(self): super(SimpleCNN, self).__init__() self.quant = quant.QuantStub() self.conv1 = nn.Conv2d(3, 16, 3, padding=1) self.relu1 = nn.ReLU() self.pool1 = nn.MaxPool2d(2) self.conv2 = nn.Conv2d(16, 32, 3, padding=1) self.relu2 = nn.ReLU() self.pool2 = nn.MaxPool2d(2) self.fc1 = nn.Linear(32 * 8 * 8, 128) self.relu3 = nn.ReLU() self.fc2 = nn.Linear(128, 10) self.dequant = quant.DeQuantStub() def forward(self, x): x = self.quant(x) x = self.conv1(x) x = self.relu1(x) x = self.pool1(x) x = self.conv2(x) x = self.relu2(x) x = self.pool2(x) x = x.view(-1, 32 * 8 * 8) x = self.fc1(x) x = self.relu3(x) x = self.fc2(x) x = self.dequant(x) return x # 初始化模型 model = SimpleCNN() # 1. 先剪枝 print("步骤1: 执行剪枝") def prune_model(model, pruning_ratio): for name, param in model.named_parameters(): if 'weight' in name: mask = torch.abs(param) > torch.quantile(torch.abs(param.view(-1)), pruning_ratio) param.data *= mask.float() return model model = prune_model(model, 0.5) # 2. 再量化 print("步骤2: 执行量化") model.eval() model.qconfig = quant.get_default_qconfig('fbgemm') quant.prepare(model, inplace=True) # 校准 calibration_data = torch.randn(100, 3, 32, 32) with torch.no_grad(): for i in range(10): batch = calibration_data[i*10:(i+1)*10] model(batch) model = quant.convert(model, inplace=True) # 计算模型大小 print(f"组合压缩后模型大小: {get_model_size(model)/1024/1024:.2f} MB") # 测试推理速度 input_data = torch.randn(1, 3, 32, 32) combined_time = measure_inference_time(model, input_data) print(f"组合压缩模型推理时间: {combined_time*1000:.2f} ms")