本文目录导读：
在Python中进行模型剪枝时保持精度是一个重要挑战,以下是几种有效方法：
渐进式剪枝（Gradual Pruning）
import torch
import torch.nn.utils.prune as prune
class GradualPruner:
def __init__(self, model, pruning_rate=0.1, target_sparsity=0.5, steps=10):
self.model = model
self.pruning_rate = pruning_rate
self.target_sparsity = target_sparsity
self.steps = steps
self.current_sparsity = 0
def step(self):
# 计算当前需要达到的稀疏度
self.current_sparsity += (self.target_sparsity - self.current_sparsity) / self.steps
for name, module in self.model.named_modules():
if isinstance(module, (torch.nn.Linear, torch.nn.Conv2d)):
# 渐进式剪枝
prune.l1_unstructured(module, name='weight', amount=self.pruning_rate)
# 移除重参数化（可选）
if self.current_sparsity >= self.target_sparsity:
for name, module in self.model.named_modules():
if isinstance(module, (torch.nn.Linear, torch.nn.Conv2d)):
prune.remove(module, 'weight')
基于重要性的剪枝（Importance-based Pruning）
class ImportanceBasedPruner:
def __init__(self, model, importance_metric='l1'):
self.model = model
self.importance_metric = importance_metric
def compute_importance(self, module):
if self.importance_metric == 'l1':
# L1范数作为重要性指标
return torch.norm(module.weight.data, p=1, dim=tuple(range(1, module.weight.dim())))
elif self.importance_metric == 'l2':
# L2范数
return torch.norm(module.weight.data, p=2, dim=tuple(range(1, module.weight.dim())))
elif self.importance_metric == 'gradient':
# 梯度幅值
return torch.norm(module.weight.grad, p=2, dim=tuple(range(1, module.weight.dim())))
def prune_by_importance(self, pruning_ratio=0.3):
for name, module in self.model.named_modules():
if isinstance(module, (torch.nn.Linear, torch.nn.Conv2d)):
importance = self.compute_importance(module)
# 保留最重要的连接
threshold = torch.quantile(importance.flatten(), pruning_ratio)
mask = importance > threshold
# 应用掩码
prune.custom_from_mask(module, name='weight', mask=mask)
知识蒸馏辅助剪枝
class DistillationPruner:
def __init__(self, teacher_model, student_model, temperature=3.0):
self.teacher = teacher_model
self.student = student_model
self.temperature = temperature
def distillation_loss(self, student_output, teacher_output, labels):
# 知识蒸馏损失
soft_targets = torch.nn.functional.softmax(
teacher_output / self.temperature, dim=1
)
soft_prob = torch.nn.functional.log_softmax(
student_output / self.temperature, dim=1
)
# 蒸馏损失 + 交叉熵损失
distill_loss = torch.nn.functional.kl_div(
soft_prob, soft_targets, reduction='batchmean'
) * (self.temperature ** 2)
ce_loss = torch.nn.functional.cross_entropy(student_output, labels)
return 0.7 * distill_loss + 0.3 * ce_loss
def train_with_distillation(self, dataloader, epochs=10):
optimizer = torch.optim.Adam(self.student.parameters(), lr=0.001)
for epoch in range(epochs):
for data, labels in dataloader:
optimizer.zero_grad()
# 教师模型不更新梯度
with torch.no_grad():
teacher_output = self.teacher(data)
student_output = self.student(data)
loss = self.distillation_loss(student_output, teacher_output, labels)
loss.backward()
optimizer.step()
完整的精度保持剪枝流程
class PrecisionAwarePruner:
def __init__(self, model, val_loader, device='cuda'):
self.model = model.to(device)
self.val_loader = val_loader
self.device = device
self.best_accuracy = 0
self.best_state_dict = None
def evaluate(self):
self.model.eval()
correct = 0
total = 0
with torch.no_grad():
for data, labels in self.val_loader:
data, labels = data.to(self.device), labels.to(self.device)
outputs = self.model(data)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
return correct / total
def iterative_prune_and_finetune(self, target_sparsity=0.5, steps=10):
current_sparsity = 0
pruner = GradualPruner(self.model, target_sparsity=target_sparsity, steps=steps)
for step in range(steps):
print(f"Step {step+1}/{steps}")
# 1. 剪枝
pruner.step()
# 2. 评估剪枝后精度
accuracy_before = self.evaluate()
print(f" Accuracy before finetune: {accuracy_before:.4f}")
# 3. 微调恢复精度
self.finetune(epochs=3)
# 4. 验证精度恢复情况
accuracy_after = self.evaluate()
print(f" Accuracy after finetune: {accuracy_after:.4f}")
# 5. 保存最佳模型
if accuracy_after > self.best_accuracy:
self.best_accuracy = accuracy_after
self.best_state_dict = self.model.state_dict().copy()
# 恢复最佳模型
if self.best_state_dict:
self.model.load_state_dict(self.best_state_dict)
print(f"Final best accuracy: {self.best_accuracy:.4f}")
return self.model
def finetune(self, train_loader=None, epochs=5, lr=0.001):
if train_loader is None:
return
self.model.train()
optimizer = torch.optim.Adam(self.model.parameters(), lr=lr)
criterion = torch.nn.CrossEntropyLoss()
for epoch in range(epochs):
for data, labels in train_loader:
data, labels = data.to(self.device), labels.to(self.device)
optimizer.zero_grad()
outputs = self.model(data)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
结构化剪枝策略
class StructuredPruner:
def __init__(self, model, pruning_ratio=0.3):
self.model = model
self.pruning_ratio = pruning_ratio
def prune_channels(self, conv_layer, pruning_ratio):
# 基于BN层权重的通道剪枝
bn_layer = None
for name, module in self.model.named_modules():
if isinstance(module, torch.nn.BatchNorm2d) and name != 'weight':
bn_layer = module
break
if bn_layer:
# 计算通道重要性
gamma = bn_layer.weight.data.abs()
sorted_gamma, _ = torch.sort(gamma)
# 找到剪枝阈值
threshold_idx = int(len(sorted_gamma) * pruning_ratio)
threshold = sorted_gamma[threshold_idx]
# 生成通道掩码
channel_mask = gamma > threshold
# 应用剪枝
conv_layer.weight.data = conv_layer.weight.data[channel_mask]
if conv_layer.bias is not None:
conv_layer.bias.data = conv_layer.bias.data[channel_mask]
# 调整BN层
bn_layer.weight.data = bn_layer.weight.data[channel_mask]
bn_layer.bias.data = bn_layer.bias.data[channel_mask]
bn_layer.running_mean = bn_layer.running_mean[channel_mask]
bn_layer.running_var = bn_layer.running_var[channel_mask]
def apply_structured_pruning(self):
for name, module in self.model.named_modules():
if isinstance(module, torch.nn.Conv2d):
self.prune_channels(module, self.pruning_ratio)
# 使用示例
if __name__ == "__main__":
# 假设有预训练模型
model = torchvision.models.resnet18(pretrained=True)
# 创建剪枝器
pruner = PrecisionAwarePruner(model, val_loader)
# 执行保持精度的剪枝
pruned_model = pruner.iterative_prune_and_finetune(
target_sparsity=0.4,
steps=8
)
关键建议：
- 渐进式剪枝：不要一次性剪太多，逐步增加稀疏度
- 微调：每次剪枝后都要进行微调
- 重要性评估：基于梯度、幅值等指标选择要保留的参数
- 知识蒸馏：利用原始模型指导剪枝后的模型学习
- 结构化剪枝：对整个通道或滤波器进行剪枝，保持计算效率
- 自动搜索：使用NAS等技术搜索最优剪枝方案
通过组合这些策略,可以在大幅减少模型参数的同时保持较高的精度。
