4. 基于梯度敏感度的结构化剪枝
任务目标
本案例针对已训练好的CNN模型，实现基于梯度敏感度的滤波器结构化剪枝：计算每个滤波器对损失函数的影响程度（梯度敏感度），剪除敏感度较低的滤波器，并在边缘设备（以树莓派4B为例）上验证推理速度提升。
核心原理
滤波器敏感度定义为：S(F_j) = (1/N) * Σ|∂L/∂W_j|，即该滤波器所有权重梯度值的均值。梯度越大，说明损失函数对该滤波器参数越敏感，删除后对精度影响越大，应保留；反之则优先剪除。结构化剪枝直接移除整个滤波器及其对应通道，无需特殊硬件支持即可获得实际加速。
环境依赖
torch==2.0.1 torchvision==0.15.1 numpy onnx onnxruntime边缘设备验证需安装onnxruntime（树莓派ARM版）。
完整实现代码
import torch
import torch.nn as nn
import torchvision.models as models
# ---------- 1. 定义带BN的ResNet-18（使用预训练权重） ----------
model = models.resnet18(pretrained=True)
model.eval()
# ---------- 2. 计算梯度敏感度 ----------
def compute_gradient_sensitivity(model, dataloader, device='cpu'):
"""计算每个Conv2d层中每个滤波器的梯度值均值"""
sensitivity = {}
grads = {}
# 注册hook收集梯度
def hook_fn(name):
def hook(module, grad_in, grad_out):
# grad_in[0] 是权重梯度
if grad_in[0] is not None:
grads[name] = grad_in[0].detach()
return hook
hooks = []
for name, module in model.named_modules():
if isinstance(module, nn.Conv2d):
hooks.append(module.register_full_backward_hook(hook_fn(name)))
sensitivity[name] = None
# 对每个batch反向传播，累积梯度
criterion = nn.CrossEntropyLoss()
for inputs, labels in dataloader:
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
loss = criterion(outputs, labels)
model.zero_grad()
loss.backward()
for name in sensitivity:
if name in grads:
# 每个滤波器梯度值均值: [out_channels]
filt_grad = grads[name].abs().mean(dim=(1, 2, 3))
if sensitivity[name] is None:
sensitivity[name] = filt_grad
else:
sensitivity[name] += filt_grad
# 取平均
for name in sensitivity:
if sensitivity[name] is not None:
sensitivity[name] /= len(dataloader)
for hook in hooks:
hook.remove()
return sensitivity
# ---------- 3. 结构化剪枝 ----------
def prune_filters(model, sensitivity, prune_ratio=0.3):
"""按敏感度剪除每个Conv层中ratio比例的滤波器"""
pruned_model = models.resnet18(pretrained=True) # 重新加载干净模型
pruned_model.eval()
# 收集所有Conv层及其后续BN层
conv_bn_pairs = []
prev_conv = None
for name, module in pruned_model.named_modules():
if isinstance(module, nn.Conv2d):
prev_conv = name
elif isinstance(module, nn.BatchNorm2d) and prev_conv is not None:
conv_bn_pairs.append((prev_conv, name))
prev_conv = None
# 对每对Conv-BN执行剪枝
for conv_name, bn_name in conv_bn_pairs:
if conv_name not in sensitivity or sensitivity[conv_name] is None:
continue
conv_module = dict(pruned_model.named_modules())[conv_name]
bn_module = dict(pruned_model.named_modules())[bn_name]
sens = sensitivity[conv_name]
# 确定保留的滤波器索引（敏感度最高的前(1-ratio)）
keep_num = int(sens.shape[0] * (1 - prune_ratio))
_, indices = torch.topk(sens, keep_num)
indices = torch.sort(indices)[0] # 升序排列
# 剪枝Conv权重: [out, in, kh, kw]
conv_module.weight.data = conv_module.weight.data[indices]
if conv_module.bias is not None:
conv_module.bias.data = conv_module.bias.data[indices]
# 剪枝BN参数
bn_module.weight.data = bn_module.weight.data[indices]
bn_module.bias.data = bn_module.bias.data[indices]
bn_module.running_mean.data = bn_module.running_mean.data[indices]
bn_module.running_var.data = bn_module.running_var.data[indices]
bn_module.num_features = keep_num
# 修改下一层Conv的输入通道数
next_conv = None
found = False
for n, m in pruned_model.named_modules():
if found and isinstance(m, nn.Conv2d):
next_conv = m
break
if n == conv_name:
found = True
if next_conv is not None:
next_conv.weight.data = next_conv.weight.data[:, indices]
return pruned_model
# ---------- 4. 导出ONNX并验证 ----------
def export_and_validate(model, dummy_input, onnx_path='pruned_model.onnx'):
torch.onnx.export(model, dummy_input, onnx_path,
input_names=['input'], output_names=['output'],
opset_version=11)
# 验证ONNX正确性
import onnx
onnx_model = onnx.load(onnx_path)
onnx.checker.check_model(onnx_model)
print(f"ONNX模型已导出并验证: {onnx_path}")
# ---------- 主流程 ----------
if __name__ == '__main__':
# 使用随机数据模拟dataloader（实际应用请替换为真实数据）
dummy_loader = [(torch.randn(4, 3, 224, 224), torch.randint(0, 1000, (4,)))]
# 计算敏感度
sens = compute_gradient_sensitivity(model, dummy_loader)
# 执行剪枝
pruned = prune_filters(model, sens, prune_ratio=0.3)
# 统计剪枝率
orig_params = sum(p.numel() for p in model.parameters())
pruned_params = sum(p.numel() for p in pruned.parameters())
print(f"原始参数量: {orig_params/1e6:.2f}M, 剪枝后: {pruned_params/1e6:.2f}M, 压缩率: {pruned_params/orig_params*100:.1f}%")
# 导出ONNX
dummy = torch.randn(1, 3, 224, 224)
export_and_validate(pruned, dummy)边缘设备部署与速度验证
# 在树莓派4B上执行（需先安装onnxruntime）
import onnxruntime as ort
import numpy as np
import time
# 加载原始与剪枝后模型
orig_session = ort.InferenceSession('original_model.onnx', providers=['CPUExecutionProvider'])
pruned_session = ort.InferenceSession('pruned_model.onnx', providers=['CPUExecutionProvider'])
# 准备输入
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
# 预热
for _ in range(5):
orig_session.run(None, {'input': input_data})
pruned_session.run(None, {'input': input_data})
# 计时（各运行50次取平均）
def benchmark(session, runs=50):
times = []
for _ in range(runs):
start = time.perf_counter()
session.run(None, {'input': input_data})
times.append(time.perf_counter() - start)
return np.mean(times) * 1000 # ms
orig_time = benchmark(orig_session)
pruned_time = benchmark(pruned_session)
print(f"原始模型推理: {orig_time:.2f}ms")
print(f"剪枝模型推理: {pruned_time:.2f}ms")
print(f"速度提升: {(1 - pruned_time/orig_time)*100:.1f}%")验证方法与通过标准
- 功能正确性：剪枝后模型能正常前向传播，输出shape不变（
[1, 1000]），ONNX检查通过。 - 结构有效性：打印的压缩率应接近
1 - (1-0.3)^2 ≈ 51%（因ResNet残差连接中部分层输入通道未剪，实际压缩率会略高）。 - 边缘加速：在树莓派4B上运行benchmark脚本，剪枝模型推理时间应明显低于原始模型（通常快30%以上）。若使用真实数据集，还需验证精度下降在可接受范围（如Top-1下降<2%）。
常见问题
- 残差连接处理：ResNet的shortcut分支与主分支相加，若shortcut的Conv层被剪枝，需同步调整相加时另一分支的通道数。本案例简化处理，仅对主路径Conv-BN对剪枝，shortcut层保持不动，因此压缩率略低于理论值。
- 敏感度计算稳定性：使用单个batch计算梯度可能噪声较大，建议使用50-100个batch取平均。若显存有限，可减小batch size。
- 剪枝后微调：本案例未包含微调步骤。实际部署前建议在训练集上微调10-20个epoch，可恢复大部分精度损失。