4. 模型剪枝

任务说明

本案例针对一个包含两个卷积层的简易CNN模型，实现通道级结构化剪枝：按L2范数评估卷积核重要性，裁剪贡献较小的输出通道，从而减少中间特征图通道数。剪枝后通过少量epoch微调恢复精度，并验证模型大小与推理速度的变化。

环境依赖

需要 torch、torchvision（仅用于MNIST数据集）。若使用CPU训练，请将下文代码中的设备改为 cpu。

完整实现

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# ---------- 1. 定义模型 ----------
class SmallCNN(nn.Module):
    def __init__(self, ch1=16, ch2=32):
        super().__init__()
        self.conv1 = nn.Conv2d(1, ch1, 3, padding=1)
        self.conv2 = nn.Conv2d(ch1, ch2, 3, padding=1)
        self.pool = nn.MaxPool2d(2)
        self.fc = nn.Linear(ch2 * 7 * 7, 10)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = self.pool(x)
        x = torch.relu(self.conv2(x))
        x = self.pool(x)
        x = x.view(x.size(0), -1)
        return self.fc(x)

# ---------- 2. 结构化剪枝函数 ----------
def prune_conv_channel(conv, ratio):
    """按L2范数裁剪conv输出通道，返回剪枝后的新卷积层"""
    l2_norm = torch.norm(conv.weight.data.view(conv.out_channels, -1), dim=1)
    keep_num = max(1, int(conv.out_channels * (1 - ratio)))
    _, indices = torch.topk(l2_norm, keep_num)
    indices = torch.sort(indices)[0]  # 保持通道顺序

    new_conv = nn.Conv2d(conv.in_channels, keep_num,
                         conv.kernel_size, conv.stride, conv.padding)
    new_conv.weight.data = conv.weight.data[indices].clone()
    new_conv.bias.data = conv.bias.data[indices].clone()
    return new_conv, indices

def structured_prune(model, ratio1, ratio2):
    """对conv1和conv2执行结构化剪枝，并调整后续层"""
    # 剪枝conv1
    new_conv1, keep_idx1 = prune_conv_channel(model.conv1, ratio1)
    # 剪枝conv2（输入通道需与conv1保留通道一致）
    old_conv2 = model.conv2
    l2_norm2 = torch.norm(old_conv2.weight.data.view(old_conv2.out_channels, -1), dim=1)
    keep_num2 = max(1, int(old_conv2.out_channels * (1 - ratio2)))
    _, idx2 = torch.topk(l2_norm2, keep_num2)
    idx2 = torch.sort(idx2)[0]

    new_conv2 = nn.Conv2d(new_conv1.out_channels, keep_num2,
                          old_conv2.kernel_size, old_conv2.stride, old_conv2.padding)
    # 同时按keep_idx1筛选输入通道，按idx2筛选输出通道
    new_conv2.weight.data = old_conv2.weight.data[idx2][:, keep_idx1].clone()
    new_conv2.bias.data = old_conv2.bias.data[idx2].clone()

    # 调整全连接层输入维度
    new_fc = nn.Linear(keep_num2 * 7 * 7, 10)
    # 全连接权重按输出通道索引重排（对应conv2保留的通道）
    fc_weight = model.fc.weight.data.view(10, old_conv2.out_channels, 7, 7)
    new_fc.weight.data = fc_weight[:, idx2].reshape(10, -1).clone()
    new_fc.bias.data = model.fc.bias.data.clone()

    pruned_model = SmallCNN(ch1=new_conv1.out_channels, ch2=new_conv2.out_channels)
    pruned_model.conv1 = new_conv1
    pruned_model.conv2 = new_conv2
    pruned_model.fc = new_fc
    return pruned_model

# ---------- 3. 训练与验证 ----------
def train(model, loader, epochs=1, lr=0.001):
    device = next(model.parameters()).device
    model.train()
    optimizer = optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()
    for _ 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
    with torch.no_grad():
        for data, target in loader:
            data, target = data.to(device), target.to(device)
            pred = model(data).argmax(dim=1)
            correct += (pred == target).sum().item()
    return correct / len(loader.dataset)

# ---------- 4. 主流程 ----------
if __name__ == "__main__":
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    transform = transforms.Compose([transforms.ToTensor(),
                                    transforms.Normalize((0.1307,), (0.3081,))])
    train_loader = DataLoader(datasets.MNIST("./data", train=True, download=True,
                                             transform=transform), batch_size=64, shuffle=True)
    test_loader = DataLoader(datasets.MNIST("./data", train=False, download=True,
                                            transform=transform), batch_size=256)

    # 训练原始模型（2个epoch作为基线）
    model = SmallCNN().to(device)
    train(model, train_loader, epochs=2)
    acc_before = evaluate(model, test_loader)
    print(f"原始模型准确率: {acc_before:.4f}")

    # 结构化剪枝：conv1保留50%，conv2保留50%
    pruned = structured_prune(model, ratio1=0.5, ratio2=0.5).to(device)
    acc_pruned = evaluate(pruned, test_loader)
    print(f"剪枝后准确率（未微调）: {acc_pruned:.4f}")

    # 微调3个epoch恢复精度
    train(pruned, train_loader, epochs=3)
    acc_finetuned = evaluate(pruned, test_loader)
    print(f"微调后准确率: {acc_finetuned:.4f}")

    # 统计参数量
    def count_params(m):
        return sum(p.numel() for p in m.parameters())
    print(f"原始参数量: {count_params(model)}")
    print(f"剪枝后参数量: {count_params(pruned)}")

运行方式

python prune_cnn.py

首次运行会自动下载MNIST数据集至 ./data 目录。

验证方法与判定标准

功能正确性：程序正常结束且无维度不匹配报错，说明剪枝后各层连接关系正确。
精度变化：观察输出——剪枝后未微调准确率应略低于原始模型（通常下降1%～5%），微调3个epoch后准确率应回升至接近原始水平（差距在2%以内视为恢复良好）。
压缩效果：参数量输出应显示剪枝后约为原始模型的1/4（因两通道均减半，全连接层同步缩小）。

若微调后精度未恢复，可适当增加微调epoch数或降低剪枝比例（如改为0.3）。

常见问题

维度不匹配：多出现在修改剪枝函数后未同步调整后续层。本实现已处理conv2输入通道索引和全连接层重排，若自行修改剪枝比例，请确保 keep_idx1 被正确用于conv2的输入筛选。
显存不足：将batch_size调小至32或16。

📚 课程目录