3. 非结构化剪枝与稀疏训练
任务目标
在PyTorch中实现细粒度（非结构化）权重剪枝：对训练后的模型按权重值阈值置零，随后进行稀疏训练（仅更新未剪枝权重）以恢复精度。最终交付可运行的剪枝与训练脚本，并给出精度对比的验证方法。
环境依赖
- Python 3.8+，PyTorch 1.12+，torchvision，NumPy
- 数据集：CIFAR-10（首次运行自动下载）
核心实现
以下代码定义了一个小型CNN，完成训练→剪枝→稀疏训练全流程。剪枝比例为50%，即保留权重值最大的50%。稀疏训练通过mask在反向传播后将已剪枝权重的梯度置零实现。
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
# ---------- 1. 定义模型 ----------
class SmallCNN(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
self.fc1 = nn.Linear(64*8*8, 256)
self.fc2 = nn.Linear(256, 10)
def forward(self, x):
x = F.relu(self.conv1(x))
x = F.max_pool2d(x, 2)
x = F.relu(self.conv2(x))
x = F.max_pool2d(x, 2)
x = x.view(x.size(0), -1)
x = F.relu(self.fc1(x))
return self.fc2(x)
# ---------- 2. 数据加载 ----------
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
])
train_ds = datasets.CIFAR10('./data', train=True, download=True, transform=transform)
test_ds = datasets.CIFAR10('./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)
# ---------- 3. 训练/测试函数 ----------
def train_epoch(model, loader, optimizer, mask=None):
model.train()
for x, y in loader:
optimizer.zero_grad()
loss = F.cross_entropy(model(x), y)
loss.backward()
if mask is not None:
for name, p in model.named_parameters():
if name in mask:
p.grad *= mask[name] # 已剪枝权重梯度置零
optimizer.step()
def evaluate(model, loader):
model.eval()
correct = total = 0
with torch.no_grad():
for x, y in loader:
pred = model(x).argmax(dim=1)
correct += (pred == y).sum().item()
total += y.size(0)
return correct / total
# ---------- 4. 剪枝函数（细粒度，按值阈值） ----------
def prune_by_threshold(model, ratio=0.5):
"""将每个权重张量中值最小的ratio比例置零，返回mask字典"""
masks = {}
for name, p in model.named_parameters():
if p.dim() == 0: # 跳过标量（如bias通常不剪）
continue
t = p.abs().view(-1)
k = int(t.numel() * ratio)
if k == 0:
continue
threshold = t.topk(t.numel() - k, largest=False).values[-1]
mask = (p.abs() > threshold).float()
p.data *= mask
masks[name] = mask
return masks
# ---------- 5. 主流程 ----------
if __name__ == '__main__':
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = SmallCNN().to(device)
# 阶段1：正常训练（5轮，演示用）
opt = optim.Adam(model.parameters(), lr=1e-3)
for _ in range(5):
train_epoch(model, train_loader, opt)
acc_before = evaluate(model, test_loader)
print(f'原始模型精度: {acc_before:.4f}')
# 阶段2：剪枝50%权重
masks = prune_by_threshold(model, ratio=0.5)
zero_ratio = sum((p==0).float().mean().item() for p in model.parameters()) / \
sum(1 for p in model.parameters())
print(f'实际置零比例: {zero_ratio:.4f}')
acc_pruned = evaluate(model, test_loader)
print(f'剪枝后精度: {acc_pruned:.4f}')
# 阶段3：稀疏训练（5轮，仅更新未剪枝权重）
opt2 = optim.Adam(model.parameters(), lr=1e-3)
for _ in range(5):
train_epoch(model, train_loader, opt2, mask=masks)
acc_final = evaluate(model, test_loader)
print(f'稀疏训练后精度: {acc_final:.4f}')
运行与验证
- 保存代码为
prune_sparse.py，执行python prune_sparse.py。 - 观察输出：
原始模型精度应约为0.60～0.70（5轮训练）；剪枝后精度会明显下降（通常低于0.50）；稀疏训练后精度应回升，且高于剪枝后精度至少3～5个百分点。 - 判定通过条件：稀疏训练后精度相比剪枝后精度有显著回升（≥3%），且最终模型权重中约50%的值为0（由
实际置零比例输出确认）。
关键说明
- 本案例使用CIFAR-10和5轮训练仅为快速演示流程；实际应用需增加训练轮数（如50轮）以获得更高基线精度。
- 剪枝阈值按每个权重张量独立计算，保证每层剪枝比例一致。
- 稀疏训练时，
mask在反向传播后对梯度置零，使已剪枝权重保持为0，未剪枝权重正常更新。 - 若需恢复更高精度，可增加稀疏训练轮数或使用学习率衰减。