📚 课程目录第3章全局阈值剪枝与局部阈值剪枝对比

3. 全局阈值剪枝与局部阈值剪枝对比

任务目标

在PyTorch中实现全局阈值剪枝与局部阈值剪枝，对比二者在相同剪枝率下的精度表现与稀疏分布特征。全局阈值对所有层统一使用一个阈值，局部阈值则按层独立设定阈值，二者均以权重值作为重要性判据。

环境依赖

Python 3.8+，PyTorch 1.12+，torchvision，NumPy。无需额外数据集，使用随机输入验证前向传播即可。

核心实现

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

class SimpleMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(28*28, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 10)
    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return self.fc3(x)

def global_threshold_prune(model, target_sparsity):
    """全局阈值：所有层共享一个阈值，达到目标稀疏度"""
    all_weights = []
    for name, param in model.named_parameters():
        if 'weight' in name:
            all_weights.append(param.data.abs().flatten())
    weights_cat = torch.cat(all_weights)
    k = int(weights_cat.numel() * target_sparsity)
    threshold = torch.kthvalue(weights_cat, k).values.item()
    total_masked = 0
    total_params = 0
    for name, param in model.named_parameters():
        if 'weight' in name:
            mask = param.data.abs() > threshold
            param.data.mul_(mask.float())
            total_masked += (~mask).sum().item()
            total_params += mask.numel()
    return threshold, total_masked / total_params

def local_threshold_prune(model, target_sparsity):
    """局部阈值：每层独立设定阈值，各层稀疏度相同"""
    thresholds = {}
    for name, param in model.named_parameters():
        if 'weight' in name:
            k = int(param.data.numel() * target_sparsity)
            thresholds[name] = torch.kthvalue(param.data.abs().flatten(), k).values.item()
            mask = param.data.abs() > thresholds[name]
            param.data.mul_(mask.float())
    return thresholds

def evaluate_accuracy(model, x, y):
    with torch.no_grad():
        out = model(x)
        pred = out.argmax(dim=1)
        return (pred == y).float().mean().item()

def layer_sparsity(model):
    result = {}
    for name, param in model.named_parameters():
        if 'weight' in name:
            result[name] = (param.data == 0).float().mean().item()
    return result

# 最小验证闭环
torch.manual_seed(42)
model = SimpleMLP()
x = torch.randn(64, 1, 28, 28)
y = torch.randint(0, 10, (64,))
target_sparsity = 0.5

# 原始精度
acc_before = evaluate_accuracy(model, x, y)

# 全局剪枝
model_global = SimpleMLP()
model_global.load_state_dict(model.state_dict())
thr_g, actual_sp_g = global_threshold_prune(model_global, target_sparsity)
acc_global = evaluate_accuracy(model_global, x, y)
sparsity_global = layer_sparsity(model_global)

# 局部剪枝
model_local = SimpleMLP()
model_local.load_state_dict(model.state_dict())
thr_l = local_threshold_prune(model_local, target_sparsity)
acc_local = evaluate_accuracy(model_local, x, y)
sparsity_local = layer_sparsity(model_local)

print(f"原始精度: {acc_before:.4f}")
print(f"全局阈值: {thr_g:.4f}, 实际稀疏度: {actual_sp_g:.4f}")
print(f"全局剪枝精度: {acc_global:.4f}")
print(f"局部剪枝精度: {acc_local:.4f}")
print("全局剪枝各层稀疏度:", {k: round(v, 4) for k, v in sparsity_global.items()})
print("局部剪枝各层稀疏度:", {k: round(v, 4) for k, v in sparsity_local.items()})
print("局部剪枝各层阈值:", {k: round(v, 4) for k, v in thr_l.items()})

运行方式

将上述代码保存为 prune_compare.py，直接执行 python prune_compare.py。代码使用随机输入数据，不依赖外部数据集，每次运行结果因随机种子固定而可复现。

验证方法与判定标准

稀疏度验证：全局剪枝的实际整体稀疏度应接近目标值0.5，但各层稀疏度差异明显——权重分布宽的层被剪更多，分布窄的层被剪更少；局部剪枝的各层稀疏度应严格接近0.5。
精度对比：在相同目标稀疏度下，局部剪枝精度通常高于全局剪枝，因为局部剪枝避免了某一层被过度剪枝导致的信息瓶颈。
阈值对比：全局剪枝只有一个阈值，局部剪枝每层一个阈值。观察各层阈值差异即可理解两种策略的本质区别。

通过条件：程序无报错运行完毕，输出中全局剪枝的实际整体稀疏度与目标值偏差小于1%，局部剪枝各层稀疏度均在0.5±0.01范围内，且精度对比结果符合上述规律。