深度學習之模型壓縮三駕馬車：基于ResNet18的模型剪枝實戰（2）

前言

《深度學習之模型壓縮三駕馬車：基于ResNet18的模型剪枝實戰（1）》里面我只是提到了對conv1層進行剪枝，只是為了驗證這個剪枝的整個過程，但是后面也有提到：僅裁剪 conv1層的影響極大，原因如下：

底層特征的重要性： conv1輸出的是最基礎的圖像特征，所有后續層的特征均基于此生成。裁剪 conv1 會直接限制后續所有層的特征表達能力。
結構連鎖反應： conv1的輸出通道減少會觸發 bn1 、 layer1.0.conv1 、 downsample 等多個模塊的調整，任何一個模塊的調整失誤（如通道數不匹配、參數初始化不當）都會導致整體性能下降。
雖然，在例子中，我們只是簡單的進行了驗證，發現效果也不是很差，但是如果具體到自己的數據，或者更加復雜的特征或者模型，可能就會影響到了整體的性能，因此，我們在原有的基礎上做了如下的改動：

剪枝目標層調整：將 conv1 改為 layer2.0.conv1 ，減少對底層特征的破壞。
通道評估優化：通過前向傳播收集激活值，優先剪枝激活值低的通道，更符合實際特征貢獻。
微調策略改進：動態解凍剪枝層及關聯的BN、downsample層，學習率降低（0.0001），微調輪次增加（10輪），確保參數充分適應。

這些修改可顯著提升剪枝后模型的穩定性和準確率。建議運行時觀察微調階段的Loss是否持續下降，若下降緩慢可進一步降低學習率（如0.00001）。
所有代碼都在這：https://gitee.com/NOON47/model_prune

詳細改動

剪枝目標層調整：將 conv1 改為 layer2.0.conv1 ，減少對底層特征的破壞。

    layer_to_prune = 'layer2.0.conv1'  # 顯式定義要剪枝的層名pruned_model = prune_conv_layer(model, layer_to_prune, amount=0.2)

通道評估優化：通過前向傳播收集激活值，優先剪枝激活值低的通道，更符合實際特征貢獻。

    model.eval()with torch.no_grad():test_input = torch.randn(128, 3, 32, 32).to(device)  # 模擬 CIFAR10 輸入features = []def hook_fn(module, input, output):features.append(output)handle = layer.register_forward_hook(hook_fn)model(test_input)handle.remove()activation = features[0]  # shape: [128, out_channels, H, W]channel_importance = activation.mean(dim=(0, 2, 3))  # 按通道求平均激活值num_channels = weight.shape[0]num_prune = int(num_channels * amount)_, indices = torch.topk(channel_importance, k=num_prune, largest=False)mask = torch.ones(num_channels, dtype=torch.bool)mask[indices] = False  # 生成剪枝掩碼

微調策略改進：動態解凍剪枝層及關聯的BN、downsample層，學習率降低（0.0001），微調輪次增加（10輪），確保參數充分適應。

    print("開始微調剪枝后的模型")# 新增：根據剪枝層動態解凍相關層（假設剪枝層為layer2.0.conv1）pruned_layer_prefix = layer_to_prune.rpartition('.')[0]  # 例如 'layer2.0'for name, param in pruned_model.named_parameters():if (pruned_layer_prefix in name) or ('fc' in name) or ('bn' in name):  # 解凍剪枝層、BN層和fc層param.requires_grad = Trueelse:param.requires_grad = Falseoptimizer = optim.Adam(filter(lambda p: p.requires_grad, pruned_model.parameters()), lr=0.0001)  # 微調學習率降低pruned_model = train_model(pruned_model, train_loader, criterion, optimizer, device, epochs=10)  # 增加微調輪次

完整的裁剪函數：

def prune_conv_layer(model, layer_name, amount=0.2):device = next(model.parameters()).devicelayer = dict(model.named_modules())[layer_name]weight = layer.weight.data# 基于激活值的通道重要性評估model.eval()with torch.no_grad():test_input = torch.randn(128, 3, 32, 32).to(device)  # 模擬 CIFAR10 輸入features = []def hook_fn(module, input, output):features.append(output)handle = layer.register_forward_hook(hook_fn)model(test_input)handle.remove()activation = features[0]  # shape: [128, out_channels, H, W]channel_importance = activation.mean(dim=(0, 2, 3))  # 按通道求平均激活值num_channels = weight.shape[0]num_prune = int(num_channels * amount)_, indices = torch.topk(channel_importance, k=num_prune, largest=False)mask = torch.ones(num_channels, dtype=torch.bool)mask[indices] = False  # 生成剪枝掩碼# 創建并替換新卷積層new_conv = nn.Conv2d(in_channels=layer.in_channels,out_channels=num_channels - num_prune,kernel_size=layer.kernel_size,stride=layer.stride,padding=layer.padding,bias=layer.bias is not None).to(device)new_conv.weight.data = layer.weight.data[mask]  # 應用掩碼剪枝權重if layer.bias is not None:new_conv.bias.data = layer.bias.data[mask]# 替換原始卷積層parent_name, sep, name = layer_name.rpartition('.')parent = model.get_submodule(parent_name)setattr(parent, name, new_conv)# 僅處理首層 conv1 的特殊邏輯if layer_name == 'conv1':# 更新首層 BN 層（bn1）bn1 = model.bn1new_bn1 = nn.BatchNorm2d(new_conv.out_channels).to(device)with torch.no_grad():new_bn1.weight.data = bn1.weight.data[mask].clone()new_bn1.bias.data = bn1.bias.data[mask].clone()new_bn1.running_mean.data = bn1.running_mean.data[mask].clone()new_bn1.running_var.data = bn1.running_var.data[mask].clone()model.bn1 = new_bn1# 處理 layer1.0 的 downsample 層（若不存在則創建）block = model.layer1[0]if not hasattr(block, 'downsample') or block.downsample is None:# 創建 1x1 卷積 + BN 用于通道匹配downsample_conv = nn.Conv2d(in_channels=new_conv.out_channels,out_channels=block.conv2.out_channels,  # 與主路徑輸出通道一致（ResNet18 為 64）kernel_size=1,stride=1,bias=False).to(device)# 初始化權重（使用原卷積層的統計量）with torch.no_grad():downsample_conv.weight.data = layer.weight.data.mean(dim=(2,3), keepdim=True)  # 原卷積核均值初始化downsample_bn = nn.BatchNorm2d(downsample_conv.out_channels).to(device)with torch.no_grad():downsample_bn.weight.data.fill_(1.0)downsample_bn.bias.data.zero_()downsample_bn.running_mean.data.zero_()downsample_bn.running_var.data.fill_(1.0)block.downsample = nn.Sequential(downsample_conv, downsample_bn)print("? 為 layer1.0 添加新的 downsample 層")else:# 調整已有 downsample 層的輸入通道downsample_conv = block.downsample[0]downsample_conv.in_channels = new_conv.out_channelsdownsample_conv.weight = nn.Parameter(downsample_conv.weight.data[:, mask, :, :].clone()).to(device)# 更新對應的 BN 層downsample_bn = block.downsample[1]new_downsample_bn = nn.BatchNorm2d(downsample_conv.out_channels).to(device)with torch.no_grad():new_downsample_bn.weight.data = downsample_bn.weight.data.clone()new_downsample_bn.bias.data = downsample_bn.bias.data.clone()new_downsample_bn.running_mean.data = downsample_bn.running_mean.data.clone()new_downsample_bn.running_var.data = downsample_bn.running_var.data.clone()block.downsample[1] = new_downsample_bn# 同步 layer1.0.conv1 的輸入通道target_conv = model.layer1[0].conv1if target_conv.in_channels != new_conv.out_channels:print(f"同步 layer1.0.conv1 輸入通道: {target_conv.in_channels} → {new_conv.out_channels}")target_conv.in_channels = new_conv.out_channelstarget_conv.weight = nn.Parameter(target_conv.weight.data[:, mask, :, :].clone()).to(device)else:# 中間層剪枝邏輯（如 layer2.0.conv1）block_prefix = layer_name.rsplit('.', 1)[0]  # 提取 block 前綴（如 'layer2.0'）block = model.get_submodule(block_prefix)     # 獲取對應的 block（如 layer2.0）# 更新當前 block 內的 BN 層（conv1 對應 bn1，conv2 對應 bn2）target_bn_name = f"{block_prefix}.bn1" if 'conv1' in layer_name else f"{block_prefix}.bn2"try:target_bn = model.get_submodule(target_bn_name)new_bn = nn.BatchNorm2d(new_conv.out_channels).to(device)with torch.no_grad():new_bn.weight.data = target_bn.weight.data[mask].clone()new_bn.bias.data = target_bn.bias.data[mask].clone()new_bn.running_mean.data = target_bn.running_mean.data[mask].clone()new_bn.running_var.data = target_bn.running_var.data[mask].clone()setattr(block, target_bn_name.split('.')[-1], new_bn)  # 替換原 BN 層print(f"? 更新剪枝層 {layer_name} 對應的 BN 層 {target_bn_name}")except AttributeError:print(f"?? 未找到剪枝層 {layer_name} 對應的 BN 層，跳過 BN 更新")# 新增：同步后續卷積層的輸入通道（如 conv1 后調整 conv2）if 'conv1' in layer_name:next_conv = block.conv2if next_conv.in_channels != new_conv.out_channels:print(f"同步 {block_prefix}.conv2 輸入通道: {next_conv.in_channels} → {new_conv.out_channels}")next_conv.in_channels = new_conv.out_channelsnext_conv.weight = nn.Parameter(next_conv.weight.data[:, mask, :, :].clone()).to(device)  # 按剪枝掩碼篩選輸入通道權重# 可選：如果存在 downsample 層，調整其輸入通道（根據實際需求啟用）# if hasattr(block, 'downsample') and block.downsample is not None:#     downsample_conv = block.downsample[0]#     downsample_conv.in_channels = new_conv.out_channels#     downsample_conv.weight = nn.Parameter(downsample_conv.weight.data[:, mask, :, :].clone()).to(device)#     print(f"? 調整剪枝層 {layer_name} 關聯的 downsample 層輸入通道")# 驗證前向傳播with torch.no_grad():test_input = torch.randn(1, 3, 32, 32).to(device)try:model(test_input)print("? 前向傳播驗證通過")except Exception as e:print(f"? 驗證失敗: {str(e)}")raisereturn model

改動后結果

經過改動后，增加微調輪次，得到的結果如下：

剪枝前模型大小信息:
==========================================================================================
Total params: 11,181,642
Trainable params: 11,181,642
Non-trainable params: 0
Total mult-adds (M): 37.03
==========================================================================================
Input size (MB): 0.01
Forward/backward pass size (MB): 0.81
Params size (MB): 44.73
Estimated Total Size (MB): 45.55
==========================================================================================
原始模型準確率: 81.42%剪枝后模型大小信息:
==========================================================================================
Total params: 11,138,392
Trainable params: 11,138,392
Non-trainable params: 0
Total mult-adds (M): 36.33
==========================================================================================
Input size (MB): 0.01
Forward/backward pass size (MB): 0.80
Params size (MB): 44.55
Estimated Total Size (MB): 45.37
==========================================================================================
剪枝后模型準確率: 83.28%