3. 全局剪枝与逐层剪枝对比

任务说明

本案例针对一个已训练好的两层全连接网络（MNIST手写数字分类），实现全局阈值剪枝：设定一个统一阈值，将网络中所有值低于该阈值的权重直接置零，然后通过重训练（fine-tune）恢复因剪枝造成的精度损失。最终对比剪枝前后及重训练后的准确率与稀疏度。

环境与依赖

需要 Python 3.8+、PyTorch 2.0+、torchvision、numpy。MNIST数据集首次运行时会自动下载。

关键完整代码

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 MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 10)
        self.relu = nn.ReLU()

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

# ---------- 2. 数据加载 ----------
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=128, shuffle=True)
test_loader = DataLoader(
    datasets.MNIST('./data', train=False, download=True, transform=transform),
    batch_size=256, shuffle=False)

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

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

# ---------- 4. 全局阈值剪枝 ----------
def global_threshold_prune(model, threshold):
    """将值低于threshold的权重置零，返回稀疏度"""
    total = 0
    zeroed = 0
    with torch.no_grad():
        for param in model.parameters():
            if param.dim() > 1:  # 只剪权重矩阵，不剪bias
                mask = param.abs() > threshold
                param.mul_(mask)
                total += mask.numel()
                zeroed += (~mask).sum().item()
    return zeroed / total

# ---------- 5. 执行剪枝与恢复 ----------
model = MLP()
print(f"初始准确率: {evaluate(model):.4f}")

train(model, epochs=2)
acc_before = evaluate(model)
print(f"训练后准确率: {acc_before:.4f}")

threshold = 0.05
sparsity = global_threshold_prune(model, threshold)
acc_pruned = evaluate(model)
print(f"剪枝后稀疏度: {sparsity:.2%}, 准确率: {acc_pruned:.4f}")

train(model, epochs=3)  # 重训练恢复精度
acc_restored = evaluate(model)
print(f"重训练后准确率: {acc_restored:.4f}")

运行与验证

直接运行上述脚本即可。判定标准：

• 剪枝后稀疏度应接近阈值对应的比例（阈值0.05时通常稀疏度在60%~80%区间）；
• 剪枝瞬间准确率会明显下降（相对训练后通常下降5~20个百分点）；
• 重训练3个epoch后准确率应回升至接近剪枝前水平（差距在2个百分点以内即视为恢复成功）。

可通过调整threshold值观察稀疏度与精度恢复的权衡关系：阈值越大，稀疏度越高，但恢复难度也越大。

常见问题

• 稀疏度为0：阈值设置过小，所有权重值均大于该值，可适当增大阈值。
• 重训练不恢复：学习率过大导致震荡，可降低学习率至1e-4并增加训练轮数。
• 只剪了部分层：本实现仅对dim()>1的参数（即权重矩阵）剪枝，bias保持不动，这是常见做法。