2. 基于L2范数的非结构化剪枝

任务目标

在MNIST分类任务上，对同一个三层全连接网络分别实施基于L1范数和L2范数的非结构化剪枝，在相同剪枝率下对比两者精度变化与权重分布差异，分析不同范数准则对剪枝结果的影响。

核心原理

非结构化剪枝逐元素移除权重。L1准则移除值最小的权重，L2准则移除平方值最小（即对输出贡献的欧氏能量最低）的权重。两者差异在于：L1对尺度均匀的小权重更敏感，L2则更倾向于保留分布集中但幅值略大的权重。通过对比剪枝后精度与剩余权重的分布形态，可观察两种准则的偏好。

完整实现

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

# ---------- 固定超参 ----------
torch.manual_seed(42)
BATCH_SIZE = 64
EPOCHS = 3
PRUNE_RATIO = 0.5  # 剪枝50%的权重

# ---------- 数据 ----------
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=BATCH_SIZE, shuffle=True)
test_loader = DataLoader(test_ds, batch_size=BATCH_SIZE, shuffle=False)

# ---------- 模型 ----------
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)

def train(model, loader, epochs=EPOCHS):
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    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

# ---------- 剪枝函数 ----------
def prune_by_norm(model, ratio, norm_type):
    """
    按范数准则剪枝：norm_type=1 使用L1，norm_type=2 使用L2
    对每个Linear层独立计算阈值并生成掩码
    """
    masks = {}
    with torch.no_grad():
        for name, param in model.named_parameters():
            if 'weight' not in name or param.dim() < 2:
                continue
            # 计算每个权重的范数贡献（L1为值，L2为平方）
            if norm_type == 1:
                scores = param.abs()
            elif norm_type == 2:
                scores = param.pow(2)
            else:
                raise ValueError("norm_type must be 1 or 2")

            # 展平后按分位数确定阈值
            flat = scores.flatten()
            k = int(flat.numel() * ratio)
            threshold = torch.kthvalue(flat, k).values

            # 生成掩码：范数得分低于阈值的置0
            mask = (scores > threshold).float()
            masks[name] = mask
            param.mul_(mask)  # 原地剪枝
    return masks

# ---------- 主流程 ----------
def run_pruning_experiment(norm_type):
    model = MLP()
    train(model, train_loader)
    acc_before = evaluate(model, test_loader)

    # 记录剪枝前权重分布
    weights_before = torch.cat([p.flatten() for n, p in model.named_parameters() if 'weight' in n and p.dim() > 1])

    # 执行剪枝
    masks = prune_by_norm(model, PRUNE_RATIO, norm_type)
    acc_after = evaluate(model, test_loader)

    # 记录剪枝后剩余权重
    weights_after = torch.cat([p.flatten() for n, p in model.named_parameters() if 'weight' in n and p.dim() > 1])

    # 统计实际剪枝比例（验证掩码生效）
    total_weights = sum(m.numel() for m in masks.values())
    zero_weights = sum((m == 0).sum().item() for m in masks.values())
    actual_ratio = zero_weights / total_weights

    return {
        'norm': norm_type,
        'acc_before': acc_before,
        'acc_after': acc_after,
        'actual_ratio': actual_ratio,
        'weights_before': weights_before,
        'weights_after': weights_after
    }

# ---------- 运行两组实验 ----------
print("=== L1 Pruning ===")
res_l1 = run_pruning_experiment(1)
print(f"Before: {res_l1['acc_before']:.4f} | After: {res_l1['acc_after']:.4f} | Actual ratio: {res_l1['actual_ratio']:.4f}")

print("\n=== L2 Pruning ===")
res_l2 = run_pruning_experiment(2)
print(f"Before: {res_l2['acc_before']:.4f} | After: {res_l2['acc_after']:.4f} | Actual ratio: {res_l2['actual_ratio']:.4f}")

# ---------- 分析：剩余权重的分布形态 ----------
def distribution_stats(weights):
    w = weights.numpy()
    return {
        'mean_abs': np.mean(np.abs(w)),
        'std': np.std(w),
        'max_abs': np.max(np.abs(w)),
        'near_zero_ratio': np.mean(np.abs(w) < 0.01)  # 接近0的权重占比
    }

print("\n=== 剪枝后剩余权重分布 ===")
print("L1 remaining:", distribution_stats(res_l1['weights_after']))
print("L2 remaining:", distribution_stats(res_l2['weights_after']))

运行方式

pip install torch torchvision numpy
python l2_pruning.py

验证方法与判定标准

剪枝生效验证：程序输出的 actual_ratio 应接近设定的0.5（允许因阈值取整有微小偏差），证明掩码正确置零了约50%的权重。
精度对比：观察两组实验的 acc_before 与 acc_after。通常L2剪枝在相同比例下精度下降略小于L1，因为L2移除的是对输出能量贡献最小的权重，而L1可能移除部分幅值小但方向重要的权重。
分布差异分析：比较两组 near_zero_ratio。L1剪枝后剩余权重中接近0的比例通常更低（因为L1已优先移除值小的），而L2剪枝后剩余权重的标准差通常更小（分布更集中），这反映两种准则对权重幅值分布的不同偏好。

预期结果参考

以下为基于固定随机种子42的确定性输出（非实测，仅作格式参考）：

=== L1 Pruning ===
Before: 0.9712 | After: 0.9543 | Actual ratio: 0.5000

=== L2 Pruning ===
Before: 0.9712 | After: 0.9618 | Actual ratio: 0.5000

=== 剪枝后剩余权重分布 ===
L1 remaining: {'mean_abs': 0.0523, 'std': 0.0612, 'max_abs': 0.3124, 'near_zero_ratio': 0.0231}
L2 remaining: {'mean_abs': 0.0487, 'std': 0.0543, 'max_abs': 0.2987, 'near_zero_ratio': 0.0412}

关键分析要点

L2剪枝精度保持更好：因为L2得分（平方）放大了大权重与中等权重的差距，阈值更容易区分“重要”与“不重要”的权重，而L1得分线性，边界附近的权重区分度较低。
L2剩余权重分布更集中（std更小）：L2倾向于保留幅值相对均匀的权重集合，而L1会保留一些幅值较大但数量更少的权重，导致分布更分散。
若两组精度差异不明显，可增大 PRUNE_RATIO 至0.7～0.9观察差异放大效应。

📚 课程目录