4. 模型剪枝

任务说明

对训练好的全连接网络实施基于权重幅度的非结构化剪枝：将值低于阈值的权重置零，随后进行微调以恢复精度。本案例使用PyTorch实现完整闭环，包含剪枝、微调与验证。

环境依赖

• Python 3.8+，PyTorch 1.12+，torchvision（仅用于MNIST数据加载）
• 无其他第三方库

完整实现

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

# ---------- 1. 模型定义 ----------
class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 10)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        return self.fc3(x)

# ---------- 2. 数据加载 ----------
def load_data(batch_size=64):
    tf = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
    train_ds = datasets.MNIST('./data', train=True, download=True, transform=tf)
    test_ds = datasets.MNIST('./data', train=False, download=True, transform=tf)
    return (DataLoader(train_ds, batch_size=batch_size, shuffle=True),
            DataLoader(test_ds, batch_size=batch_size))

# ---------- 3. 训练/微调与评估 ----------
def train(model, loader, epochs=1, lr=0.001):
    optimizer = optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()
    model.train()
    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
    total = 0
    with torch.no_grad():
        for data, target in loader:
            pred = model(data).argmax(dim=1)
            correct += (pred == target).sum().item()
            total += target.size(0)
    return correct / total

# ---------- 4. 权重剪枝 ----------
def prune_weights(model, threshold):
    """将值低于threshold的权重置零，返回剪枝比例"""
    total = 0
    pruned = 0
    for param in model.parameters():
        if param.dim() < 2:  # 跳过bias
            continue
        mask = param.abs() > threshold
        total += param.numel()
        pruned += (~mask).sum().item()
        param.data.mul_(mask.float())
    return pruned / total

# ---------- 5. 主流程 ----------
if __name__ == '__main__':
    train_loader, test_loader = load_data()

    # 训练基线模型
    model = MLP()
    train(model, train_loader, epochs=3, lr=0.001)
    base_acc = evaluate(model, test_loader)
    print(f'基线准确率: {base_acc:.4f}')

    # 剪枝
    threshold = 0.02
    ratio = prune_weights(model, threshold)
    pruned_acc = evaluate(model, test_loader)
    print(f'剪枝比例: {ratio:.2%}, 剪枝后准确率: {pruned_acc:.4f}')

    # 微调恢复
    train(model, train_loader, epochs=2, lr=0.0005)
    finetuned_acc = evaluate(model, test_loader)
    print(f'微调后准确率: {finetuned_acc:.4f}')

运行方式

python prune_mnist.py

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

验证方法

• 观察控制台输出：基线准确率通常约0.97；剪枝后准确率下降；微调后准确率回升至接近基线水平。
• 判定通过条件：剪枝比例大于50%（阈值0.02时通常可达70%以上），且微调后准确率不低于基线2个百分点以内。
• 可通过调整threshold观察剪枝比例与精度权衡：阈值越大，剪枝越多，精度损失越大，微调恢复越困难。

关键说明

• 剪枝仅作用于权重矩阵（dim >= 2），偏置不参与剪枝，避免破坏网络结构。
• 非结构化剪枝产生稀疏权重，但不会减少存储或计算量（需专用稀疏库支持），本案例聚焦于剪枝方法本身。
• 微调使用更低学习率（0.0005），避免破坏已保留权重的有效信息。