1. 基于PyTorch的ResNet-18结构化剪枝
任务说明
对ResNet-18实施通道剪枝（结构化剪枝），在CIFAR-10数据集上完成训练、剪枝与微调，对比剪枝前后模型的Top-1精度与FLOPs变化。本案例聚焦于卷积层输出通道的剪枝流程，不涉及BN层折叠与推理加速内核实现。
环境依赖
- Python 3.8+，PyTorch 1.12+，torchvision 0.13+
- numpy，tqdm（可选）
实现代码
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import numpy as np
# ---------- 1. 数据加载 ----------
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2)
testloader = torch.utils.data.DataLoader(testset, batch_size=256, shuffle=False, num_workers=2)
# ---------- 2. 模型定义 ----------
def conv3x3(in_planes, out_planes, stride=1):
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)
class BasicBlock(nn.Module):
expansion = 1
def __init__(self, inplanes, planes, stride=1, downsample=None):
super().__init__()
self.conv1 = conv3x3(inplanes, planes, stride)
self.bn1 = nn.BatchNorm2d(planes)
self.relu = nn.ReLU(inplace=True)
self.conv2 = conv3x3(planes, planes)
self.bn2 = nn.BatchNorm2d(planes)
self.downsample = downsample
self.stride = stride
def forward(self, x):
identity = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
if self.downsample is not None:
identity = self.downsample(x)
out += identity
out = self.relu(out)
return out
class ResNet18(nn.Module):
def __init__(self, num_classes=10):
super().__init__()
self.inplanes = 64
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
self.bn1 = nn.BatchNorm2d(64)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = self._make_layer(BasicBlock, 64, 2)
self.layer2 = self._make_layer(BasicBlock, 128, 2, stride=2)
self.layer3 = self._make_layer(BasicBlock, 256, 2, stride=2)
self.layer4 = self._make_layer(BasicBlock, 512, 2, stride=2)
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.fc = nn.Linear(512, num_classes)
def _make_layer(self, block, planes, blocks, stride=1):
downsample = None
if stride != 1 or self.inplanes != planes * block.expansion:
downsample = nn.Sequential(
nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(planes * block.expansion),
)
layers = []
layers.append(block(self.inplanes, planes, stride, downsample))
self.inplanes = planes * block.expansion
for _ in range(1, blocks):
layers.append(block(self.inplanes, planes))
return nn.Sequential(*layers)
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = self.avgpool(x)
x = torch.flatten(x, 1)
x = self.fc(x)
return x
# ---------- 3. 训练/测试工具 ----------
def train(model, epochs=1, lr=0.01):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=5e-4)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
model.train()
for epoch in range(epochs):
for inputs, targets in trainloader:
inputs, targets = inputs.to(device), targets.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
scheduler.step()
def evaluate(model):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
model.eval()
correct = 0
total = 0
with torch.no_grad():
for inputs, targets in testloader:
inputs, targets = inputs.to(device), targets.to(device)
outputs = model(inputs)
_, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
return 100.0 * correct / total
# ---------- 4. 通道剪枝实现 ----------
def prune_resnet18_channels(model, prune_ratio=0.3):
"""按比例剪枝每个BasicBlock中conv2的输出通道（即残差块内部第二层卷积）。"""
pruned_model = ResNet18(num_classes=10)
pruned_model.load_state_dict(model.state_dict())
for name, module in pruned_model.named_modules():
if isinstance(module, BasicBlock):
# 获取conv2的BN层gamma值
bn2 = module.bn2
gamma = bn2.weight.detach().cpu().numpy()
num_channels = gamma.shape[0]
k = int(num_channels * prune_ratio)
if k == 0:
continue
# 选择gamma最小的k个通道索引
indices = np.argsort(np.abs(gamma))[:k]
mask = torch.ones(num_channels, dtype=torch.bool)
mask[indices] = False
# 剪枝conv2输出通道
conv2 = module.conv2
new_conv2 = nn.Conv2d(conv2.in_channels, num_channels - k,
kernel_size=conv2.kernel_size, stride=conv2.stride,
padding=conv2.padding, bias=False)
keep_idx = torch.where(mask)[0]
with torch.no_grad():
new_conv2.weight.data = conv2.weight.data[keep_idx]
module.conv2 = new_conv2
# 剪枝bn2
new_bn2 = nn.BatchNorm2d(num_channels - k)
with torch.no_grad():
new_bn2.weight.data = bn2.weight.data[keep_idx]
new_bn2.bias.data = bn2.bias.data[keep_idx]
new_bn2.running_mean.data = bn2.running_mean.data[keep_idx]
new_bn2.running_var.data = bn2.running_var.data[keep_idx]
module.bn2 = new_bn2
# 调整下一个卷积层的输入通道（若存在）
# 找到该block之后的下一个卷积层
parent_name = name
# 处理layer内部的顺序：block内conv2之后是残差连接，下一个卷积在下一层或下个block
# 简化处理：遍历找到下一个使用该block输出的卷积层
# 此处通过模块遍历实现
next_conv = None
found = False
for n, m in pruned_model.named_modules():
if found and isinstance(m, nn.Conv2d):
next_conv = m
break
if n == parent_name:
found = True
if next_conv is not None:
old_weight = next_conv.weight.data
new_next_conv = nn.Conv2d(num_channels - k, next_conv.out_channels,
kernel_size=next_conv.kernel_size, stride=next_conv.stride,
padding=next_conv.padding, bias=next_conv.bias is not None)
with torch.no_grad():
new_next_conv.weight.data = old_weight[:, keep_idx]
if next_conv.bias is not None:
new_next_conv.bias.data = next_conv.bias.data
# 替换
for n, m in pruned_model.named_modules():
if m is next_conv:
parent = pruned_model
parts = n.split('.')
for p in parts[:-1]:
parent = getattr(parent, p)
setattr(parent, parts[-1], new_next_conv)
break
return pruned_model
# ---------- 5. FLOPs计算 ----------
def compute_flops(model, input_size=(1, 3, 32, 32)):
"""简化FLOPs计算：仅统计卷积层和全连接层乘法次数。"""
model.eval()
flops = 0
hooks = []
def conv_hook(module, input, output):
nonlocal flops
batch, in_c, h, w = input[0].size()
out_c, _, kh, kw = module.weight.size()
out_h, out_w = output.size(2), output.size(3)
flops += batch * out_c * out_h * out_w * in_c * kh * kw
def fc_hook(module, input, output):
nonlocal flops
flops += input[0].size(0) * input[0].size(1) * output.size(1)
for m in model.modules():
if isinstance(m, nn.Conv2d):
hooks.append(m.register_forward_hook(conv_hook))
elif isinstance(m, nn.Linear):
hooks.append(m.register_forward_hook(fc_hook))
dummy = torch.zeros(input_size)
model(dummy)
for h in hooks:
h.remove()
return flops / 1e6 # 返回MFLOPs
# ---------- 6. 主流程 ----------
if __name__ == '__main__':
# 训练原始模型（为节省时间，可加载预训练权重或减少epoch）
print("训练原始ResNet-18...")
model = ResNet18(num_classes=10)
# 实际使用时建议加载预训练权重或训练更久
train(model, epochs=5, lr=0.01)
acc_orig = evaluate(model)
flops_orig = compute_flops(model)
print(f"原始模型: Top-1 Acc = {acc_orig:.2f}%, FLOPs = {flops_orig:.2f}M")
# 剪枝
print("执行通道剪枝...")
pruned_model = prune_resnet18_channels(model, prune_ratio=0.3)
acc_pruned = evaluate(pruned_model)
flops_pruned = compute_flops(pruned_model)
print(f"剪枝后(未微调): Top-1 Acc = {acc_pruned:.2f}%, FLOPs = {flops_pruned:.2f}M")
# 微调
print("微调剪枝模型...")
train(pruned_model, epochs=5, lr=0.005)
acc_finetuned = evaluate(pruned_model)
print(f"微调后: Top-1 Acc = {acc_finetuned:.2f}%, FLOPs = {flops_finetuned:.2f}M")
print(f"FLOPs 减少: {(1 - flops_pruned/flops_orig)*100:.1f}%")运行方式
- 将上述代码保存为
resnet18_prune.py，首次运行会自动下载CIFAR-10数据集（约170MB）至./data目录。 - 执行
python resnet18_prune.py。若GPU可用则自动使用GPU，否则回退CPU（训练时间会显著增加）。 - 为快速验证流程，示例中训练轮数设为5；实际使用建议加载torchvision预训练权重或训练50轮以上以获得基准精度。
验证方法
- 观察终端输出：原始模型精度应接近CIFAR-10合理水平（训练充分时约90%+，5轮快速验证约60-75%）。
- 剪枝后未微调精度会明显下降（通常下降5-15个百分点），微调后精度回升。
- FLOPs输出应显示剪枝后计算量减少约30%（与prune_ratio对应），验证剪枝生效。
- 通过条件：程序无异常执行完毕，输出三个阶段的精度与FLOPs数值，且剪枝后FLOPs低于原始模型。
常见问题
- 维度不匹配错误：剪枝后残差连接的维度需保持一致。本实现仅剪枝BasicBlock内部conv2的通道，不影响shortcut分支，因此不会出现维度冲突。
- 精度过低：5轮训练仅为流程演示，实际使用请增加训练轮数或加载预训练权重。
- FLOPs计算精度：此处为简化统计，未包含BN层和ReLU的FLOPs，但卷积与全连接层占主导，相对变化趋势准确。