4. 全局阈值通道裁剪与模型微调

任务目标

在边缘AI部署中，模型参数量直接影响存储与推理延迟。本案例以ResNet-18的BN层gamma值为通道重要性指标，通过设定全局阈值裁剪低贡献通道，并对裁剪后的模型进行微调以恢复精度。核心路径为：统计gamma分布→确定阈值→通道裁剪→微调验证。

环境与依赖

torch==2.0.1
torchvision==0.15.2
numpy==1.24.3

关键实现

以下代码实现完整的裁剪与微调闭环。核心逻辑：遍历所有BN层，将gamma值低于全局阈值的通道标记为待裁剪；通过构建新卷积层实现通道数缩减；微调仅需标准训练循环。

import torch
import torch.nn as nn
import torchvision.models as models
from torch.utils.data import DataLoader, TensorDataset

def analyze_bn_gamma(model):
    """收集所有BN层gamma值，返回分布统计"""
    gammas = []
    for name, module in model.named_modules():
        if isinstance(module, nn.BatchNorm2d):
            gammas.extend(module.weight.detach().cpu().numpy().tolist())
    return np.array(gammas)

def get_pruning_mask(model, threshold):
    """根据全局阈值生成每层通道保留掩码"""
    masks = {}
    for name, module in model.named_modules():
        if isinstance(module, nn.BatchNorm2d):
            gamma = module.weight.detach().abs()
            masks[name] = gamma > threshold
    return masks

def prune_conv_bn_pair(conv, bn, mask):
    """裁剪单个Conv-BN对，返回新conv和bn"""
    idx = torch.where(mask)[0]
    # 裁剪卷积核
    new_conv = nn.Conv2d(
        conv.in_channels, len(idx),
        conv.kernel_size, conv.stride,
        conv.padding, conv.dilation, conv.groups, conv.bias is not None
    )
    new_conv.weight.data = conv.weight.data[idx]
    if conv.bias is not None:
        new_conv.bias.data = conv.bias.data[idx]
    # 裁剪BN
    new_bn = nn.BatchNorm2d(len(idx))
    new_bn.weight.data = bn.weight.data[idx]
    new_bn.bias.data = bn.bias.data[idx]
    new_bn.running_mean.data = bn.running_mean.data[idx]
    new_bn.running_var.data = bn.running_var.data[idx]
    return new_conv, new_bn

def prune_model(model, threshold):
    """主裁剪函数：处理所有Conv-BN对"""
    # 先收集所有需要裁剪的模块对
    modules_to_prune = []
    prev_name = None
    for name, module in model.named_modules():
        if isinstance(module, nn.Conv2d):
            prev_name = name
        elif isinstance(module, nn.BatchNorm2d) and prev_name:
            modules_to_prune.append((prev_name, name))
            prev_name = None
    
    # 从后往前裁剪，避免索引错乱
    for conv_name, bn_name in reversed(modules_to_prune):
        conv = dict(model.named_modules())[conv_name]
        bn = dict(model.named_modules())[bn_name]
        mask = get_pruning_mask(model, threshold)[bn_name]
        if mask.sum() == 0:
            continue
        new_conv, new_bn = prune_conv_bn_pair(conv, bn, mask)
        # 替换模块
        parent_conv = model
        for part in conv_name.split('.')[:-1]:
            parent_conv = getattr(parent_conv, part)
        setattr(parent_conv, conv_name.split('.')[-1], new_conv)
        parent_bn = model
        for part in bn_name.split('.')[:-1]:
            parent_bn = getattr(parent_bn, part)
        setattr(parent_bn, bn_name.split('.')[-1], new_bn)
    return model

def finetune(model, train_loader, epochs=3, lr=0.001):
    """标准微调流程"""
    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)
    model.train()
    for epoch in range(epochs):
        total_loss = 0
        for inputs, labels in train_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        print(f'Epoch {epoch+1}, Loss: {total_loss/len(train_loader):.4f}')

# ========== 使用示例 ==========
if __name__ == '__main__':
    # 1. 加载预训练模型
    model = models.resnet18(pretrained=True)
    
    # 2. 分析gamma分布
    gammas = analyze_bn_gamma(model)
    print(f'Gamma统计: mean={gammas.mean():.4f}, std={gammas.std():.4f}, p50={np.percentile(gammas, 50):.4f}')
    
    # 3. 设定阈值（示例：取gamma分布的20分位数）
    threshold = np.percentile(gammas, 20)
    print(f'裁剪阈值: {threshold:.4f}')
    
    # 4. 执行裁剪
    pruned_model = prune_model(model, threshold)
    
    # 5. 统计裁剪率
    orig_params = sum(p.numel() for p in model.parameters())
    pruned_params = sum(p.numel() for p in pruned_model.parameters())
    print(f'参数裁剪率: {(1 - pruned_params/orig_params)*100:.1f}%')
    
    # 6. 微调（使用随机数据演示流程）
    dummy_input = torch.randn(64, 3, 224, 224)
    dummy_label = torch.randint(0, 1000, (64,))
    loader = DataLoader(TensorDataset(dummy_input, dummy_label), batch_size=16)
    finetune(pruned_model, loader, epochs=1)
    
    # 7. 验证前向传播
    pruned_model.eval()
    with torch.no_grad():
        test_out = pruned_model(torch.randn(1, 3, 224, 224))
    print(f'输出维度: {test_out.shape}')

运行与验证

运行方式：直接执行上述脚本。首次运行会下载ResNet-18预训练权重（约45MB）。

验证路径：

1. 观察gamma分布统计输出，确认阈值设定合理（通常取10%~30%分位数）
2. 检查参数裁剪率输出，预期在15%~35%之间（取决于阈值）
3. 确认微调loss逐epoch下降，说明模型正在恢复精度
4. 最终前向传播输出维度应为[1, 1000]，证明模型结构完整

通过条件：程序无报错完成全流程，裁剪率与阈值设定一致，微调后loss低于初始值。若需验证精度恢复效果，可在真实数据集（如CIFAR-10）上对比裁剪前后top-1准确率。

常见问题

• 维度不匹配：裁剪后下一层输入通道数变化，需确保所有Conv-Bn对都被处理。本实现通过modules_to_prune列表保证成对处理。
• 阈值过高导致某层全裁：代码中已添加mask.sum()==0跳过保护，但建议阈值不超过gamma分布的30分位数。
• 微调收敛慢：可适当增大学习率至0.005，或增加微调轮数至5轮。