2. 基于L1范数的通道剪枝
任务目标
对训练好的卷积神经网络执行基于L1范数的通道剪枝：计算每个卷积层各通道权重值和，按预设比例剪除L1范数最小的通道，重建剪枝后的模型结构，并通过短时微调恢复精度。本案例以CIFAR-10上的简易CNN为例，完整演示“训练→剪枝→微调→验证”闭环。
环境依赖
Python 3.8+，PyTorch 1.12+，torchvision，numpy。无需额外硬件，CPU即可完成本案例全部流程。
实现原理与关键代码
通道剪枝的核心在于：对卷积层输出通道的权重矩阵（形状为 [out_channels, in_channels, k, k]）计算每个输出通道的L1范数（即该通道全部权重值之和），该值越小说明该通道对后续特征贡献越弱，可安全移除。剪枝后需同步删除下一层对应输入通道，并重建网络结构。
以下代码实现完整流程：先训练一个基准CNN，再执行L1通道剪枝（剪除比例50%），最后微调5个epoch恢复精度。
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
# ---------- 1. 定义简易CNN ----------
class SimpleCNN(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
self.conv3 = nn.Conv2d(32, 64, 3, padding=1)
self.pool = nn.MaxPool2d(2)
self.fc = nn.Linear(64 * 4 * 4, 10)
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = self.pool(F.relu(self.conv3(x)))
x = x.view(x.size(0), -1)
return self.fc(x)
# ---------- 2. 数据加载 ----------
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
train_set = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
test_set = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
train_loader = DataLoader(train_set, batch_size=64, shuffle=True)
test_loader = DataLoader(test_set, batch_size=256)
# ---------- 3. 训练/测试工具函数 ----------
def train(model, loader, epochs=1, lr=0.001):
device = next(model.parameters()).device
optimizer = optim.Adam(model.parameters(), lr=lr)
criterion = nn.CrossEntropyLoss()
model.train()
for epoch in range(epochs):
for data, target in loader:
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
loss = criterion(model(data), target)
loss.backward()
optimizer.step()
def evaluate(model, loader):
device = next(model.parameters()).device
model.eval()
correct = 0
total = 0
with torch.no_grad():
for data, target in loader:
data, target = data.to(device), target.to(device)
output = model(data)
pred = output.argmax(dim=1)
correct += (pred == target).sum().item()
total += target.size(0)
return correct / total
# ---------- 4. L1范数通道剪枝 ----------
def prune_channels(model, prune_ratio=0.5):
"""按L1范数剪除每个卷积层prune_ratio比例的通道，返回剪枝后的新模型"""
# 收集所有卷积层
conv_layers = [m for m in model.modules() if isinstance(m, nn.Conv2d)]
# 记录每层保留的通道索引
keep_indices = []
for conv in conv_layers:
# 计算每个输出通道的L1范数
l1_norm = conv.weight.abs().sum(dim=(1, 2, 3)) # [out_channels]
num_keep = int(l1_norm.size(0) * (1 - prune_ratio))
# 保留L1范数最大的num_keep个通道
_, indices = torch.topk(l1_norm, num_keep)
keep_indices.append(indices.sort()[0]) # 升序排列，保持通道顺序
# 重建模型
new_model = SimpleCNN()
# 逐层替换卷积权重和偏置
new_convs = [m for m in new_model.modules() if isinstance(m, nn.Conv2d)]
old_convs = conv_layers
# conv1: 输入通道不变，输出通道取keep_indices[0]
idx0 = keep_indices[0]
new_convs[0].weight.data = old_convs[0].weight.data[idx0].clone()
new_convs[0].bias.data = old_convs[0].bias.data[idx0].clone()
# conv2: 输入通道对应conv1保留的通道，输出通道取keep_indices[1]
idx1 = keep_indices[1]
new_convs[1].weight.data = old_convs[1].weight.data[idx1][:, idx0].clone()
new_convs[1].bias.data = old_convs[1].bias.data[idx1].clone()
# conv3: 输入通道对应conv2保留的通道，输出通道取keep_indices[2]
idx2 = keep_indices[2]
new_convs[2].weight.data = old_convs[2].weight.data[idx2][:, idx1].clone()
new_convs[2].bias.data = old_convs[2].bias.data[idx2].clone()
# 全连接层：输入维度需匹配conv3输出通道数 * 4 * 4
new_fc_in = idx2.size(0) * 4 * 4
new_model.fc = nn.Linear(new_fc_in, 10)
# 从原fc中截取对应权重（简单平均池化映射，此处用前new_fc_in行）
with torch.no_grad():
new_model.fc.weight.data = model.fc.weight.data[:, :new_fc_in].clone()
new_model.fc.bias.data = model.fc.bias.data.clone()
return new_model
# ---------- 5. 主流程 ----------
if __name__ == '__main__':
# 训练基准模型（为节省时间仅训练3个epoch，实际应用可训练更久）
model = SimpleCNN()
print("训练基准模型...")
train(model, train_loader, epochs=3)
acc_before = evaluate(model, test_loader)
print(f"剪枝前准确率: {acc_before:.4f}")
# 统计剪枝前参数量
params_before = sum(p.numel() for p in model.parameters())
# 执行剪枝
pruned_model = prune_channels(model, prune_ratio=0.5)
params_after = sum(p.numel() for p in pruned_model.parameters())
print(f"参数量: {params_before} -> {params_after} (减少{100*(1-params_after/params_before):.1f}%)")
# 剪枝后直接评估（未微调）
acc_pruned = evaluate(pruned_model, test_loader)
print(f"剪枝后(未微调)准确率: {acc_pruned:.4f}")
# 微调恢复精度
print("微调中...")
train(pruned_model, train_loader, epochs=5, lr=0.0005)
acc_finetuned = evaluate(pruned_model, test_loader)
print(f"微调后准确率: {acc_finetuned:.4f}")
运行方式与验证
直接运行上述脚本即可。首次运行会自动下载CIFAR-10数据集（约170MB）。程序依次输出：剪枝前准确率、参数量变化、剪枝后未微调准确率、微调后准确率。
通过条件：
- 剪枝后参数量减少约50%（因全连接层占比较大，实际总参数量减少约40%~50%属正常）；
- 剪枝后未微调准确率较剪枝前下降（通常下降5~15个百分点）；
- 微调5个epoch后准确率回升，与剪枝前差距缩小到3个百分点以内即视为剪枝+微调流程有效。
关键说明
- 本案例中全连接层权重采用直接截断方式适配，严格做法应基于全局平均池化或特征图重计算，但对本演示目标（验证L1剪枝有效性）已足够；
- 剪枝比例0.5为示例值，实际工程中建议从0.3开始逐步调高，观察精度-参数量权衡曲线；
- 若需部署加速，需配合硬件支持的稀疏卷积库或实际通道移除后的模型导出，本案例聚焦于剪枝算法本身。