4. 全局剪枝与局部剪枝对比

任务目标

在PyTorch中实现两种结构化剪枝策略——全局阈值剪枝（所有层共享一个稀疏度阈值）与逐层剪枝（每层按各自稀疏度独立裁剪），在相同总稀疏度下对比二者对模型精度的影响，并输出压缩率与精度变化。

环境与依赖

Python 3.8+，PyTorch 1.12+，torchvision，NumPy。无需额外硬件，CPU即可完成本实验。

实现思路

使用LeNet-5在MNIST上训练3个epoch作为基线，随后分别执行全局剪枝与逐层剪枝（均为非结构化细粒度剪枝），最后微调1个epoch并评估。全局剪枝将所有卷积/全连接层的权重值合并排序，按全局阈值置零；逐层剪枝则对每层单独排序，使每层稀疏度一致。两种方法保持总稀疏度相同（50%），以便公平对比。

关键完整代码

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

# ---------- 模型定义 ----------
class LeNet5(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 6, 5, padding=2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16*5*5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)
    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), 2)
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return self.fc3(x)

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

# ---------- 训练/评估工具 ----------
def train(model, epochs=1):
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    model.train()
    for _ in range(epochs):
        for x, y in train_loader:
            opt.zero_grad()
            loss = F.cross_entropy(model(x), y)
            loss.backward()
            opt.step()

def evaluate(model):
    model.eval()
    correct = total = 0
    with torch.no_grad():
        for x, y in test_loader:
            pred = model(x).argmax(dim=1)
            correct += (pred == y).sum().item()
            total += y.size(0)
    return correct / total

def count_zeros(model):
    total = zeros = 0
    for p in model.parameters():
        if p.dim() > 1:  # 只统计权重，不统计bias
            total += p.numel()
            zeros += (p == 0).sum().item()
    return zeros / total

# ---------- 剪枝实现 ----------
def global_prune(model, sparsity):
    """全局阈值剪枝：所有层合并排序后统一置零"""
    weights = [p for p in model.parameters() if p.dim() > 1]
    all_vals = torch.cat([w.abs().flatten() for w in weights])
    threshold = torch.quantile(all_vals, sparsity)
    with torch.no_grad():
        for w in weights:
            w.masked_fill_(w.abs() < threshold, 0)

def layerwise_prune(model, sparsity):
    """逐层剪枝：每层独立按相同稀疏度裁剪"""
    with torch.no_grad():
        for p in model.parameters():
            if p.dim() > 1:
                threshold = torch.quantile(p.abs().flatten(), sparsity)
                p.masked_fill_(p.abs() < threshold, 0)

# ---------- 主流程 ----------
print("=== 训练基线模型 ===")
model_base = LeNet5()
train(model_base, epochs=3)
acc_base = evaluate(model_base)
print(f"基线精度: {acc_base:.4f}")

# 全局剪枝
model_global = LeNet5()
model_global.load_state_dict(model_base.state_dict())
global_prune(model_global, sparsity=0.5)
train(model_global, epochs=1)
acc_global = evaluate(model_global)
sp_global = count_zeros(model_global)

# 逐层剪枝
model_layer = LeNet5()
model_layer.load_state_dict(model_base.state_dict())
layerwise_prune(model_layer, sparsity=0.5)
train(model_layer, epochs=1)
acc_layer = evaluate(model_layer)
sp_layer = count_zeros(model_layer)

# ---------- 结果输出 ----------
def calc_compression(model):
    """计算压缩率：非零参数占比的倒数"""
    total = nonzero = 0
    for p in model.parameters():
        if p.dim() > 1:
            total += p.numel()
            nonzero += (p != 0).sum().item()
    return total / max(nonzero, 1)

print("\n=== 剪枝结果对比 ===")
print(f"全局剪枝: 实际稀疏度={sp_global:.4f}, 压缩率={calc_compression(model_global):.2f}x, 精度={acc_global:.4f}")
print(f"逐层剪枝: 实际稀疏度={sp_layer:.4f}, 压缩率={calc_compression(model_layer):.2f}x, 精度={acc_layer:.4f}")
print(f"精度差(全局-逐层): {acc_global - acc_layer:+.4f}")

运行方式

将上述代码保存为prune_compare.py，在项目目录下执行：

python prune_compare.py

首次运行会自动下载MNIST数据集（约10MB），需保持网络畅通。

验证方法与判定标准

程序输出应包含三部分：基线精度、两种剪枝策略的稀疏度/压缩率/精度。通过标准如下：

• 两种剪枝的实际稀疏度均接近0.5（误差<0.01），证明剪枝量一致；
• 全局剪枝与逐层剪枝的精度差在±0.02以内属正常（MNIST上两者差异通常较小）；
• 若逐层剪枝精度显著高于全局剪枝（>0.03），说明该模型各层冗余度不均，逐层策略更优；反之则全局策略更优。

可修改sparsity参数（如0.3、0.7）重复实验，观察不同稀疏度下两种策略的精度变化趋势。

常见问题

• 稀疏度不准确：torch.quantile在数据量小时可能有微小偏差，若需精确稀疏度可改用排序后取第k个值作为阈值。
• 精度大幅下降：当稀疏度>0.7时，LeNet-5精度会明显退化，建议在0.3~0.6范围内对比。
• bias未剪枝：代码仅对dim>1的权重剪枝，bias保持稠密，这是结构化剪枝的常见做法，不影响对比公平性。