2. 模型量化基础

任务说明

对给定的简易CNN模型执行通道剪枝：按BN层缩放因子γ的L1范数排序，裁剪贡献较低的通道，从而减少模型参数与FLOPs，并通过短时微调恢复精度。本案例聚焦剪枝流程本身，使用CIFAR-10的一个固定子集（5000张训练图/1000张测试图）以保证可复现性。

环境依赖

Python 3.8+，PyTorch 1.12+，torchvision 0.13+。无需GPU，CPU即可在数分钟内完成全部流程。

关键完整代码

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import Subset
import torchvision
import torchvision.transforms as transforms
import numpy as np

# ---------- 1. 数据准备（固定子集，保证可复现） ----------
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
train_subset = Subset(trainset, list(range(5000)))
test_subset = Subset(testset, list(range(1000)))
train_loader = torch.utils.data.DataLoader(train_subset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_subset, batch_size=64, shuffle=False)

# ---------- 2. 简易CNN（含BN层，用于通道剪枝） ----------
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
        self.bn3 = nn.BatchNorm2d(128)
        self.pool = nn.MaxPool2d(2)
        self.fc = nn.Linear(128 * 4 * 4, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.bn1(self.conv1(x))))
        x = self.pool(F.relu(self.bn2(self.conv2(x))))
        x = self.pool(F.relu(self.bn3(self.conv3(x))))
        x = x.view(x.size(0), -1)
        return self.fc(x)

# ---------- 3. 训练/测试工具 ----------
def train(model, loader, epochs=3):
    model.train()
    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
    criterion = nn.CrossEntropyLoss()
    for _ in range(epochs):
        for data, target in loader:
            optimizer.zero_grad()
            loss = criterion(model(data), target)
            loss.backward()
            optimizer.step()

def evaluate(model, loader):
    model.eval()
    correct = 0
    with torch.no_grad():
        for data, target in loader:
            pred = model(data).argmax(dim=1)
            correct += (pred == target).sum().item()
    return correct / len(loader.dataset)

# ---------- 4. 通道剪枝实现 ----------
def prune_channels(model, ratio=0.3):
    """按BN层gamma值剪枝，返回剪枝后的新模型"""
    # 收集所有BN层的gamma
    bn_layers = [m for m in model.modules() if isinstance(m, nn.BatchNorm2d)]
    all_gamma = torch.cat([bn.weight.data.abs().view(-1) for bn in bn_layers])
    k = int(len(all_gamma) * ratio)
    threshold = all_gamma.sort()[0][k]  # 第k小的gamma值作为阈值

    # 为每层生成保留通道掩码
    masks = []
    for bn in bn_layers:
        masks.append(bn.weight.data.abs() > threshold)

    # 构建新模型
    new_model = SimpleCNN()
    new_state = new_model.state_dict()

    # 复制未剪枝层的参数（fc层直接复制）
    old_state = model.state_dict()
    for name in new_state:
        if 'fc' in name:
            new_state[name] = old_state[name]

    # 逐层处理卷积+BN
    conv_bn_pairs = [
        ('conv1', 'bn1', 'conv2'),
        ('conv2', 'bn2', 'conv3'),
        ('conv3', 'bn3', 'fc')
    ]
    prev_mask = torch.ones(3, dtype=torch.bool)  # 输入通道全保留

    for conv_name, bn_name, next_name in conv_bn_pairs:
        conv = getattr(model, conv_name)
        bn = getattr(model, bn_name)
        mask = masks[list(bn_layers).index(bn)]

        # 剪枝当前卷积的输出通道（即BN的通道）
        new_conv = getattr(new_model, conv_name)
        new_bn = getattr(new_model, bn_name)

        # 保留的通道索引
        keep_idx = torch.nonzero(mask).view(-1)

        # 更新卷积权重: [out_keep, in_keep, kh, kw]
        w = conv.weight.data[mask][:, prev_mask]
        new_conv.weight.data.copy_(w)
        if conv.bias is not None:
            new_conv.bias.data.copy_(conv.bias.data[mask])

        # 更新BN参数
        new_bn.weight.data.copy_(bn.weight.data[mask])
        new_bn.bias.data.copy_(bn.bias.data[mask])
        new_bn.running_mean.copy_(bn.running_mean[mask])
        new_bn.running_var.copy_(bn.running_var[mask])

        # 更新下一层的输入通道
        if next_name != 'fc':
            next_conv = getattr(model, next_name)
            new_next_conv = getattr(new_model, next_name)
            # 下一层输入通道对应当前层输出通道
            new_next_conv.weight.data.copy_(
                next_conv.weight.data[:, mask, :, :]
            )
        else:
            # fc层输入维度需要调整
            fc = model.fc
            new_fc = new_model.fc
            # 展平后维度 = 128*4*4，需按mask调整
            # 此处简化：仅当最后一层被剪枝时调整fc
            if not mask.all():
                # 计算保留通道数
                keep = mask.sum().item()
                # 重新初始化fc并复制对应权重（需知道展平顺序）
                # 展平顺序为 [C, H, W]，此处C=128, H=W=4
                # 权重形状 [10, 128*16]，需按通道维切分
                w_fc = fc.weight.data.view(10, 128, 4, 4)
                new_w = w_fc[:, mask, :, :].reshape(10, -1)
                new_fc.weight.data.copy_(new_w)
                new_fc.bias.data.copy_(fc.bias.data)

        prev_mask = mask

    return new_model

# ---------- 5. 主流程 ----------
if __name__ == '__main__':
    torch.manual_seed(42)
    np.random.seed(42)

    # 训练原始模型
    model = SimpleCNN()
    print("训练原始模型...")
    train(model, train_loader, epochs=3)
    acc_before = evaluate(model, test_loader)
    print(f"原始模型测试准确率: {acc_before:.4f}")

    # 统计原始参数量
    orig_params = sum(p.numel() for p in model.parameters())
    print(f"原始参数量: {orig_params}")

    # 执行剪枝
    pruned_model = prune_channels(model, ratio=0.3)
    pruned_params = sum(p.numel() for p in pruned_model.parameters())
    print(f"剪枝后参数量: {pruned_params}，减少 {100*(1-pruned_params/orig_params):.1f}%")

    # 剪枝后直接评估（未微调）
    acc_pruned = evaluate(pruned_model, test_loader)
    print(f"剪枝后未微调准确率: {acc_pruned:.4f}")

    # 微调恢复精度
    print("微调剪枝模型...")
    train(pruned_model, train_loader, epochs=3)
    acc_finetuned = evaluate(pruned_model, test_loader)
    print(f"微调后准确率: {acc_finetuned:.4f}")

运行方式

python channel_pruning.py

首次运行会自动下载CIFAR-10数据集（约170MB）到./data目录。

验证方法与通过条件

• 剪枝有效性：控制台输出中“剪枝后参数量”应明显小于“原始参数量”（约减少30%），证明通道确实被裁剪。
• 精度恢复：微调后的准确率应高于剪枝后未微调的准确率，且接近原始模型准确率（通常差距在2%以内即视为有效恢复）。
• 代码自洽性：程序无报错完整跑通，输出三个准确率数值，且满足上述趋势关系。

常见问题

• 最后一层剪枝后fc维度不匹配：本代码已处理该情况，通过将fc权重reshape为[10, C, 4, 4]后按mask裁剪再展平。
• 剪枝后模型无法加载：由于直接构建新模型并复制参数，需确保新旧模型结构一致（均为SimpleCNN类），仅通道数不同。
• 阈值选择：本实现采用全局阈值（所有BN层gamma混合排序），若需按层独立剪枝可修改prune_channels中的排序逻辑。