1. 基于PyTorch的LeNet-5模型剪枝与量化
任务目标
在MNIST数据集上完成LeNet-5的完整流程：训练基线模型 → 结构化剪枝（按卷积通道剪枝）→ 后训练量化（PTQ，8位动态量化）→ 对比剪枝/量化前后的模型大小与测试精度。
环境依赖
Python 3.8+，PyTorch 1.13+，torchvision，numpy。MNIST数据集首次运行自动下载（需网络或提前放置./data目录）。
关键完整代码
import torch
import torch.nn as nn
import torch.nn.utils.prune as prune
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
# ---------- 1. 定义LeNet-5（适配MNIST 28x28输入） ----------
class LeNet5(nn.Module):
def __init__(self):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(1, 6, 5, padding=2), nn.ReLU(), nn.MaxPool2d(2), # 6@14x14
nn.Conv2d(6, 16, 5), nn.ReLU(), nn.MaxPool2d(2), # 16@5x5
)
self.classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(16*5*5, 120), nn.ReLU(),
nn.Linear(120, 84), nn.ReLU(),
nn.Linear(84, 10)
)
def forward(self, x):
return self.classifier(self.features(x))
# ---------- 2. 训练/测试工具 ----------
def train(model, loader, epochs=1, lr=0.001):
device = next(model.parameters()).device
model.train()
opt = torch.optim.Adam(model.parameters(), lr=lr)
loss_fn = nn.CrossEntropyLoss()
for _ in range(epochs):
for x, y in loader:
x, y = x.to(device), y.to(device)
opt.zero_grad()
loss = loss_fn(model(x), y)
loss.backward()
opt.step()
def evaluate(model, loader):
device = next(model.parameters()).device
model.eval()
correct = total = 0
with torch.no_grad():
for x, y in loader:
x, y = x.to(device), y.to(device)
pred = model(x).argmax(dim=1)
correct += (pred == y).sum().item()
total += y.size(0)
return correct / total
def model_size_mb(model):
return sum(p.numel()*p.element_size() for p in model.parameters()) / 1e6
# ---------- 3. 结构化剪枝：按卷积输出通道全局阈值 ----------
def structured_prune(model, amount=0.3):
"""对两个卷积层按L2范数剪掉指定比例的整个输出通道"""
convs = [m for m in model.features if isinstance(m, nn.Conv2d)]
# 收集所有卷积核的L2范数
norms = []
for conv in convs:
w = conv.weight.detach() # [out_c, in_c, kh, kw]
norms.append(w.view(w.size(0), -1).norm(dim=1))
all_norms = torch.cat(norms)
threshold = torch.quantile(all_norms, amount)
for conv in convs:
w = conv.weight.detach()
keep = w.view(w.size(0), -1).norm(dim=1) > threshold
# 若全部被剪则保留范数最大的一个，避免通道数为0
if not keep.any():
keep[w.view(w.size(0), -1).norm(dim=1).argmax()] = True
# 用prune方法将不保留通道的权重置零（结构化掩码）
prune.custom_from_mask(conv, name='weight', mask=keep.view(-1, 1, 1, 1).expand_as(w).float())
prune.remove(conv, 'weight') # 固化掩码，实际删除参数
# ---------- 4. 主流程 ----------
def main():
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
tf = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
train_loader = DataLoader(datasets.MNIST('./data', train=True, download=True, transform=tf), batch_size=64, shuffle=True)
test_loader = DataLoader(datasets.MNIST('./data', train=False, download=True, transform=tf), batch_size=256)
# 训练基线
model = LeNet5().to(device)
train(model, train_loader, epochs=2)
acc_base = evaluate(model, test_loader)
size_base = model_size_mb(model)
# 结构化剪枝（30%通道）
structured_prune(model, amount=0.3)
# 剪枝后微调1个epoch恢复精度
train(model, train_loader, epochs=1, lr=0.0005)
acc_pruned = evaluate(model, test_loader)
size_pruned = model_size_mb(model)
# PTQ：8位动态量化（仅量化Linear层）
model_quant = torch.quantization.quantize_dynamic(model, {nn.Linear}, dtype=torch.qint8)
acc_quant = evaluate(model_quant, test_loader)
size_quant = model_size_mb(model_quant)
print(f'基线: 精度 {acc_base:.4f}, 模型 {size_base:.3f} MB')
print(f'剪枝后: 精度 {acc_pruned:.4f}, 模型 {size_pruned:.3f} MB')
print(f'量化后: 精度 {acc_quant:.4f}, 模型 {size_quant:.3f} MB')
if __name__ == '__main__':
main()运行方式
python lenet_prune_quant.py验证方法与通过条件
- 精度对比：剪枝后精度相对基线下降通常<1%，量化后精度下降通常<0.5%（MNIST任务较简单，实际波动以运行为准）。
- 模型大小：量化后模型大小应显著小于基线（动态量化主要压缩Linear层，LeNet-5中Linear层占参数量约97%，预期可减小至约1/4）。
- 功能正确性：程序无报错跑完，输出三行精度与大小数据即视为闭环完成。
关键说明
- 结构化剪枝实现为通道级置零并固化掩码，实际参数量未减少但权重稀疏化；如需真正减少计算量，需配合稀疏卷积后端或导出ONNX后处理。
- 动态量化仅量化Linear层，Conv层保持FP32，因此精度损失较小。
- MNIST训练2轮精度约99%左右，剪枝后微调1轮可恢复大部分精度损失。