2. 基于BN层缩放因子的通道剪枝
任务目标
本案例实现基于BN层缩放因子的通道剪枝：以ResNet-18每个BN层的γ缩放因子作为通道重要性指标，对γ值较小的通道进行裁剪，在ImageNet子集（100类）上对比剪枝前后模型的Top-1准确率与参数量，并分析不同剪枝比例的影响。
核心原理
BN层对每个通道执行 y = γ·x̂ + β，其中γ的值大小反映了该通道对输出的贡献程度。训练时对γ施加L1正则化使其稀疏化，剪枝时移除γ接近0的通道，即可在保持精度的前提下压缩模型。
环境依赖
torch==2.0.1 torchvision==0.15.2 numpy tqdm timm==0.9.2关键完整代码
import torch
import torch.nn as nn
import torchvision.models as models
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, Subset
import numpy as np
import copy
# ---------- 1. 稀疏化训练（带L1正则的BN γ） ----------
def train_sparse(model, loader, epochs=3, lr=1e-3, sparsity=1e-4):
optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=5e-4)
criterion = nn.CrossEntropyLoss()
model.train()
for epoch in range(epochs):
for x, y in loader:
optimizer.zero_grad()
loss = criterion(model(x), y)
# 对BN层γ施加L1正则
l1 = sum(p.abs().sum() for name, p in model.named_parameters() if 'bn.weight' in name)
(loss + sparsity * l1).backward()
optimizer.step()
# ---------- 2. 通道剪枝 ----------
def prune_channels(model, ratio):
"""按BN层γ值排序，裁剪最小ratio比例的通道"""
model = copy.deepcopy(model)
bn_modules = [m for m in model.modules() if isinstance(m, nn.BatchNorm2d)]
keep_idx = {}
for name, m in model.named_modules():
if isinstance(m, nn.BatchNorm2d):
gamma = m.weight.detach().abs().cpu().numpy()
k = int(len(gamma) * (1 - ratio))
idx = np.argsort(gamma)[-k:] # 保留γ最大的k个通道
keep_idx[name] = torch.tensor(np.sort(idx))
return model, keep_idx
# ---------- 3. 实际裁剪（重建卷积层） ----------
def apply_pruning(model, keep_idx):
"""根据keep_idx重建所有Conv2d和BN层"""
for name, module in model.named_modules():
if isinstance(module, nn.Conv2d):
parent_name = name.rsplit('.', 1)[0]
bn_name = parent_name + '.bn'
if bn_name in keep_idx:
idx = keep_idx[bn_name]
# 裁剪输入通道
if module.in_channels > 1:
module.in_channels = len(idx)
module.weight = nn.Parameter(module.weight.data[:, idx, :, :])
# 裁剪输出通道（由下一个BN决定）
next_bn = dict(model.named_modules()).get(bn_name)
if next_bn is not None:
out_idx = keep_idx[bn_name]
module.out_channels = len(out_idx)
module.weight = nn.Parameter(module.weight.data[out_idx, :, :, :])
if module.bias is not None:
module.bias = nn.Parameter(module.bias.data[out_idx])
return model
# ---------- 4. 验证与统计 ----------
def evaluate(model, loader):
model.eval()
correct = 0
total = 0
with torch.no_grad():
for x, y in loader:
out = model(x)
pred = out.argmax(dim=1)
correct += (pred == y).sum().item()
total += y.size(0)
return correct / total
def count_params(model):
return sum(p.numel() for p in model.parameters())
# ---------- 5. 主流程 ----------
def main():
# 数据准备（ImageNet子集100类）
transform = transforms.Compose([
transforms.Resize(224),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
# 假设数据位于./imagenet_subset，含train/和val/子目录
train_ds = datasets.ImageFolder('./imagenet_subset/train', transform)
val_ds = datasets.ImageFolder('./imagenet_subset/val', transform)
train_loader = DataLoader(train_ds, batch_size=64, shuffle=True, num_workers=4)
val_loader = DataLoader(val_ds, batch_size=64, shuffle=False, num_workers=4)
# 加载预训练ResNet-18
model = models.resnet18(pretrained=True)
model.fc = nn.Linear(512, 100) # 适配100类
# 稀疏化训练
print("开始稀疏化训练...")
train_sparse(model, train_loader, epochs=3)
# 评估原始模型
acc_orig = evaluate(model, val_loader)
params_orig = count_params(model)
print(f"原始模型: Acc={acc_orig:.4f}, Params={params_orig/1e6:.2f}M")
# 不同剪枝比例实验
for ratio in [0.3, 0.5, 0.7]:
pruned_model, keep_idx = prune_channels(model, ratio)
pruned_model = apply_pruning(pruned_model, keep_idx)
acc_pruned = evaluate(pruned_model, val_loader)
params_pruned = count_params(pruned_model)
print(f"剪枝比例{ratio:.0%}: Acc={acc_pruned:.4f}, Params={params_pruned/1e6:.2f}M, "
f"压缩率={(1-params_pruned/params_orig)*100:.1f}%")
if __name__ == '__main__':
main()运行方式
# 1. 准备ImageNet子集（100类，train/val目录结构）
# 2. 执行主程序
python channel_pruning.py验证方法与判定标准
- 功能验证：程序能完整运行，输出原始模型与各剪枝比例的准确率和参数量。
- 有效性判定：剪枝30%时准确率下降应小于2个百分点；剪枝50%时下降小于5个百分点；参数量压缩率与剪枝比例基本一致。
- 正确性检查：打印的参数量随剪枝比例增大而递减，且模型能正常前向传播（无维度错误）。
常见问题
- 维度不匹配：若剪枝后报错，检查shortcut连接处的卷积层是否也需同步裁剪（ResNet残差分支需额外处理）。
- 准确率骤降：稀疏化训练epoch不足导致γ未充分稀疏化，可增加epoch或调大sparsity系数。
- 数据准备：ImageNet子集需按标准目录结构放置，否则ImageFolder无法正确加载。
说明：图中曲线为预期趋势示意，实际数值取决于数据子集与训练轮次。剪枝比例增大时，准确率缓慢下降而压缩率线性上升，在50%剪枝处通常能取得较好的精度-压缩平衡。