4. 分层阈值剪枝

任务目标

实现分层阈值剪枝：为模型每一层独立计算剪枝阈值并移除低于该阈值的权重，与全局统一阈值剪枝进行稀疏度分布与精度对比。

核心实现

import torch
import torch.nn as nn
import torch.nn.utils.prune as prune
from torchvision.models import resnet18

def layerwise_prune(model, sparsity):
    """按层独立阈值剪枝：每层保留sparsity比例的最大权重"""
    for name, module in model.named_modules():
        if isinstance(module, nn.Conv2d) or isinstance(module, nn.Linear):
            # 计算该层独立阈值（按分位数）
            tensor = module.weight.data
            threshold = torch.quantile(tensor.abs().flatten(), sparsity)
            # 生成掩码：值大于阈值的保留
            mask = tensor.abs() > threshold
            prune.custom_from_mask(module, name='weight', mask=mask)
    return model

def global_threshold_prune(model, sparsity):
    """全局统一阈值：所有层共享一个阈值"""
    all_weights = []
    for module in model.modules():
        if isinstance(module, (nn.Conv2d, nn.Linear)):
            all_weights.append(module.weight.data.abs().flatten())
    all_weights = torch.cat(all_weights)
    threshold = torch.quantile(all_weights, sparsity)
    
    for module in model.modules():
        if isinstance(module, (nn.Conv2d, nn.Linear)):
            mask = module.weight.data.abs() > threshold
            prune.custom_from_mask(module, name='weight', mask=mask)
    return model

def measure_sparsity(model):
    """统计每层实际稀疏度"""
    stats = {}
    for name, module in model.named_modules():
        if isinstance(module, (nn.Conv2d, nn.Linear)):
            mask = module.weight_mask
            sparsity = 1 - mask.float().mean().item()
            stats[name] = sparsity
    return stats

# 使用示例
model = resnet18(pretrained=False)
# 分层剪枝：每层剪掉50%最小权重
layerwise_prune(model, sparsity=0.5)
layer_stats = measure_sparsity(model)

# 重置模型做全局剪枝对比
model2 = resnet18(pretrained=False)
global_threshold_prune(model2, sparsity=0.5)
global_stats = measure_sparsity(model2)

# 输出对比（示例输出，非实测）
print("分层剪枝各层稀疏度范围: {:.2f} ~ {:.2f}".format(
    min(layer_stats.values()), max(layer_stats.values())))
print("全局剪枝各层稀疏度范围: {:.2f} ~ {:.2f}".format(
    min(global_stats.values()), max(global_stats.values())))

运行与验证

1. 依赖：PyTorch ≥1.9，torchvision（仅用于获取ResNet-18结构）。
2. 执行：直接运行上述脚本，观察两种策略下各层稀疏度分布。
3. 通过条件：
   - 分层剪枝后，各层稀疏度应接近设定的50%（因权重分布不同略有波动）；
   - 全局剪枝后，权重分布较宽的层稀疏度偏高，较窄的层稀疏度偏低，层间差异明显大于分层剪枝；
   - 若需精度验证，可在剪枝后调用torch.optim微调模型，对比两种方法在验证集上的Top-1准确率。

关键说明

分层阈值剪枝的核心在于每层独立计算阈值，避免全局阈值对权重分布不均匀层造成过度剪枝或欠剪枝。实际应用中，通常对敏感层设置较低稀疏度、对冗余层设置较高稀疏度，而非统一50%。上述代码使用torch.nn.utils.prune实现掩码剪枝，剪枝后模型仍保留原始结构，可通过model.weight_mask查看掩码。