1. PyTorch模型剪枝基础

任务目标

使用PyTorch实现一个简单CNN模型，对其卷积层执行权重剪枝（将值小于阈值的权重置零），并对比剪枝前后模型的参数量、文件大小与测试集准确率，验证剪枝在模型压缩与精度保持之间的权衡。

环境依赖

• Python 3.8+
• PyTorch 1.12+（含torchvision）
• NumPy

完整实现

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
import numpy as np
import os, copy

# ---------- 1. 定义简单CNN ----------
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 8, 3, padding=1)   # 输入1通道，输出8通道
        self.conv2 = nn.Conv2d(8, 16, 3, padding=1)
        self.fc1 = nn.Linear(16*7*7, 64)
        self.fc2 = nn.Linear(64, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))   # 14x14
        x = F.relu(F.max_pool2d(self.conv2(x), 2))   # 7x7
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        return self.fc2(x)

# ---------- 2. 数据加载 ----------
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('./data', train=True, download=True, transform=transform),
    batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(
    datasets.MNIST('./data', train=False, download=True, transform=transform),
    batch_size=256, shuffle=False)

# ---------- 3. 训练函数 ----------
def train(model, epochs=2):
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    model.train()
    for epoch in range(epochs):
        for data, target in train_loader:
            optimizer.zero_grad()
            loss = F.cross_entropy(model(data), target)
            loss.backward()
            optimizer.step()

# ---------- 4. 评估函数 ----------
def evaluate(model):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for data, target in test_loader:
            out = model(data)
            pred = out.argmax(dim=1)
            correct += (pred == target).sum().item()
            total += target.size(0)
    return 100.0 * correct / total

# ---------- 5. 权重剪枝（全局阈值） ----------
def prune_weights(model, threshold):
    """将值小于threshold的权重置零，返回剪枝比例"""
    pruned_model = copy.deepcopy(model)
    total = 0
    zeroed = 0
    for name, param in pruned_model.named_parameters():
        if 'weight' in name and param.dim() >= 2:  # 只剪卷积/全连接权重，不剪bias
            mask = torch.abs(param.data) > threshold
            zeroed += (mask == 0).sum().item()
            total += param.numel()
            param.data *= mask.float()
    ratio = zeroed / total if total > 0 else 0.0
    return pruned_model, ratio

# ---------- 6. 模型大小统计 ----------
def model_size_mb(model):
    torch.save(model.state_dict(), './tmp_model.pth')
    size = os.path.getsize('./tmp_model.pth') / 1024 / 1024
    os.remove('./tmp_model.pth')
    return size

# ---------- 7. 主流程 ----------
if __name__ == '__main__':
    # 训练原始模型
    model = SimpleCNN()
    train(model, epochs=2)
    acc_before = evaluate(model)
    size_before = model_size_mb(model)

    # 执行剪枝（阈值设为0.01）
    pruned_model, ratio = prune_weights(model, threshold=0.01)
    acc_after = evaluate(pruned_model)
    size_after = model_size_mb(pruned_model)

    # 输出结果
    print(f'原始模型: 准确率={acc_before:.2f}%, 大小={size_before:.4f}MB')
    print(f'剪枝模型: 准确率={acc_after:.2f}%, 大小={size_after:.4f}MB')
    print(f'权重剪枝比例: {ratio*100:.1f}%')
    print(f'大小缩减: {(1-size_after/size_before)*100:.1f}%')

运行方式

pip install torch torchvision numpy
python prune_cnn.py

验证方法与判定标准

• 功能验证：程序正常输出四行统计信息，无报错。
• 剪枝生效：输出的“权重剪枝比例”应大于0%（阈值0.01下通常可剪去30%~70%的权重）。
• 压缩效果：“大小缩减”应大于0%，证明剪枝后模型文件确实变小。
• 精度保持：剪枝后准确率与剪枝前相比下降幅度通常在5个百分点以内（MNIST简单任务上，若阈值过大可能下降更多，可调小阈值至0.005验证）。

关键说明

• 本案例采用全局阈值剪枝，即对所有卷积/全连接层权重统一按值阈值置零，未做逐层敏感度分析或重训练（fine-tune），属于基础剪枝演示。
• 模型大小统计基于state_dict保存后的文件体积，稀疏权重在标准格式下仍按稠密存储，因此大小缩减主要来自被剪枝权重数值变为0后可能带来的存储优化（实际PyTorch默认格式下缩减有限，若需更大压缩需配合稀疏存储或量化）。
• 若准确率下降过多，可降低阈值（如0.005）或增加训练轮次后再剪枝。