4. 训练后全整型量化实战
任务目标
本案例基于L1范数通道剪枝方法，对CIFAR-10分类模型进行结构化剪枝。通过计算卷积层每个通道权重张量的L1范数作为重要性指标，裁剪贡献较小的通道，实现模型压缩，并对比剪枝前后的模型参数量与精度变化。
环境依赖
Python 3.8+，PyTorch 1.10+，torchvision，NumPy。无需额外硬件，CPU即可完成本案例全部流程。
实现原理与代码
核心思路：对每个卷积层，计算其权重张量在输出通道维度上的L1范数（即每个输出通道对应卷积核的值之和），按比例保留范数较大的通道，裁剪其余通道。剪枝后需重建相邻层的输入通道维度，保持网络结构合法。
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms
import numpy as np
# ---------- 1. 定义简单CNN模型 ----------
class SimpleCNN(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
self.conv3 = nn.Conv2d(32, 64, 3, padding=1)
self.pool = nn.MaxPool2d(2)
self.fc = nn.Linear(64 * 4 * 4, 10)
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = self.pool(F.relu(self.conv3(x)))
x = x.view(x.size(0), -1)
return self.fc(x)
# ---------- 2. L1范数通道剪枝函数 ----------
def prune_conv_layer(layer, next_layer, keep_ratio):
"""按L1范数裁剪卷积层输出通道，并调整下一层输入通道"""
# 计算每个输出通道的L1范数
l1_norm = layer.weight.abs().sum(dim=(1, 2, 3))
num_keep = max(int(l1_norm.size(0) * keep_ratio), 1)
_, indices = torch.topk(l1_norm, num_keep)
indices = torch.sort(indices).values
# 裁剪当前层
layer.weight.data = layer.weight.data[indices]
layer.bias.data = layer.bias.data[indices]
layer.out_channels = num_keep
# 调整下一层（若存在）
if next_layer is not None:
if isinstance(next_layer, nn.Conv2d):
next_layer.weight.data = next_layer.weight.data[:, indices]
next_layer.in_channels = num_keep
elif isinstance(next_layer, nn.Linear):
# 需先计算池化后的特征图尺寸
next_layer.weight.data = next_layer.weight.data.view(
next_layer.out_features, -1, 4, 4)[:, indices].view(
next_layer.out_features, -1)
next_layer.in_features = num_keep * 4 * 4
# ---------- 3. 训练/测试工具 ----------
def train(model, loader, epochs=1):
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()
model.train()
for epoch in range(epochs):
for x, y in loader:
optimizer.zero_grad()
loss = criterion(model(x), y)
loss.backward()
optimizer.step()
def evaluate(model, loader):
model.eval()
correct = 0
total = 0
with torch.no_grad():
for x, y in loader:
out = model(x)
pred = out.argmax(dim=1)
correct += (pred == y).sum().item()
total += y.size(0)
return correct / total
def count_params(model):
return sum(p.numel() for p in model.parameters())
# ---------- 4. 主流程 ----------
if __name__ == "__main__":
# 数据准备（仅用少量数据演示流程）
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
download=True, transform=transform)
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(testset, batch_size=64)
# 初始化并训练原始模型
model = SimpleCNN()
print("原始模型参数量:", count_params(model))
train(model, train_loader, epochs=1) # 演示用1轮，实际可增加
acc_before = evaluate(model, test_loader)
print(f"剪枝前精度: {acc_before:.4f}")
# 执行剪枝（保留50%通道）
prune_conv_layer(model.conv1, model.conv2, keep_ratio=0.5)
prune_conv_layer(model.conv2, model.conv3, keep_ratio=0.5)
prune_conv_layer(model.conv3, model.fc, keep_ratio=0.5)
print("剪枝后模型参数量:", count_params(model))
acc_after = evaluate(model, test_loader)
print(f"剪枝后精度（未微调）: {acc_after:.4f}")
# 微调恢复精度
train(model, train_loader, epochs=1)
acc_finetuned = evaluate(model, test_loader)
print(f"微调后精度: {acc_finetuned:.4f}")运行与验证
直接执行脚本即可完成全流程。首次运行会自动下载CIFAR-10数据集（约170MB）。验证通过条件：
- 程序无异常退出，输出四行指标信息；
- 剪枝后参数量显著小于原始参数量（约减少50%以上）；
- 剪枝后精度下降在可接受范围（通常<5%），微调后精度回升。
关键说明
本案例为教学演示，训练仅1轮且未做数据增强，因此精度值偏低属正常现象。实际应用中应增加训练轮数（如10-20轮）以获得更高基线精度。剪枝比例可通过keep_ratio参数调节，本案例统一使用0.5。若需对全连接层之前的卷积层剪枝，需注意view操作中的特征图尺寸需与实际网络结构匹配。