PyTorch实战：用Network Slimming给ResNet模型‘瘦身’，推理速度提升30%

在边缘计算设备上部署深度学习模型时，模型大小和推理速度往往是制约因素。最近在Jetson Nano上部署ResNet-34时，我发现原模型不仅占用大量内存，推理速度也难以满足实时性要求。经过多次尝试，Network Slimming这种基于BN层权重的剪枝方法，最终让模型体积缩小了65%，推理速度提升30%以上，而精度损失控制在1%以内。下面分享完整实现过程。

1. 环境准备与基准测试

1.1 硬件与软件配置

实验使用Jetson Nano开发板（4GB内存版）作为测试平台，关键配置如下：

组件	规格
CPU	四核Cortex-A57 @ 1.43GHz
GPU	128核Maxwell架构
内存	4GB LPDDR4
存储	32GB eMMC 5.1
系统	Ubuntu 18.04 LTS
PyTorch	1.10.0 with JetPack 4.6

安装必要的Python包：

pip install torch==1.10.0 torchvision==0.11.1 numpy==1.19.5

1.2 基准模型测试

使用torchvision提供的预训练ResNet-34在CIFAR-10上进行微调后，测试原始模型性能：

import torch
from torchvision.models import resnet34

model = resnet34(pretrained=True)
# 修改最后一层适配CIFAR-10的10分类
model.fc = torch.nn.Linear(512, 10) 

# 测试推理速度
input_tensor = torch.randn(1, 3, 32, 32).cuda()
with torch.no_grad():
    for _ in range(100):
        _ = model(input_tensor)

测试结果：

• 模型大小：85.3MB
• 单张图片推理时间：23.4ms
• 测试集准确率：94.7%

2. Network Slimming实现详解

2.1 稀疏化训练关键代码

Network Slimming的核心是在训练过程中对BN层权重施加L1正则：

def update_bn_weights(model, s=0.0001):
    """BN层稀疏化更新"""
    for module in model.modules():
        if isinstance(module, torch.nn.BatchNorm2d):
            module.weight.grad.data.add_(
                s * torch.sign(module.weight.data)
            )

在训练循环中调用：

optimizer.step()
update_bn_weights(model)  # 在每个batch后执行

提示：稀疏系数s需要谨慎选择，过大会导致权重过度稀疏，过小则剪枝效果不明显。建议从1e-4开始尝试。

2.2 自动剪枝阈值计算

训练完成后，我们需要确定各层的剪枝比例。这里采用全局排序法：

def calculate_threshold(model, percent=0.5):
    """计算全局剪枝阈值"""
    bn_weights = []
    for m in model.modules():
        if isinstance(m, torch.nn.BatchNorm2d):
            bn_weights.append(m.weight.data.abs().clone())
    
    all_weights = torch.cat(bn_weights)
    threshold = torch.quantile(all_weights, percent)
    return threshold

2.3 通道剪枝实现

根据阈值生成剪枝掩码并应用到模型：

def prune_model(model, threshold):
    masks = []
    for m in model.modules():
        if isinstance(m, torch.nn.BatchNorm2d):
            mask = (m.weight.data.abs() > threshold).float()
            masks.append(mask)
            # 应用掩码
            m.weight.data.mul_(mask)
            m.bias.data.mul_(mask)
    return masks

3. 剪枝后模型重构

3.1 模型结构自动调整

剪枝后需要重建一个更紧凑的模型：

def rebuild_model(old_model, masks):
    new_cfg = [int(torch.sum(mask)) for mask in masks]
    
    # 根据new_cfg构建新的ResNet
    new_model = ResNet(block=BasicBlock, layers=[3,4,6,3], cfg=new_cfg)
    
    # 复制保留的权重
    old_modules = list(old_model.modules())
    new_modules = list(new_model.modules())
    
    layer_id = 0
    for m_old, m_new in zip(old_modules, new_modules):
        if isinstance(m_old, torch.nn.BatchNorm2d):
            mask = masks[layer_id]
            idx = torch.nonzero(mask).squeeze()
            # 复制BN层参数
            m_new.weight.data = m_old.weight.data[idx].clone()
            m_new.bias.data = m_old.bias.data[idx].clone()
            layer_id += 1
            
        elif isinstance(m_old, torch.nn.Conv2d):
            # 处理卷积层...
            pass
            
    return new_model

3.2 微调策略

剪枝后的模型需要经过微调恢复精度：

• 学习率策略：初始学习率设为原训练时的1/10
• 数据增强：增加MixUp、CutMix等增强方法
• 训练周期：通常需要原训练时间的1/3到1/2

optimizer = torch.optim.SGD(
    pruned_model.parameters(), 
    lr=0.001,  # 初始学习率
    momentum=0.9,
    weight_decay=1e-4
)

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer, T_max=50
)

4. 效果对比与部署优化

4.1 性能指标对比

经过50%剪枝比例后的结果：

指标	原始模型	剪枝后模型	变化率
参数量	21.3M	7.2M	↓66.2%
模型大小	85.3MB	28.9MB	↓66.1%
推理时延	23.4ms	16.1ms	↓31.2%
准确率	94.7%	94.1%	↓0.6%

4.2 部署优化技巧

在Jetson Nano上进一步优化的方法：

1. TensorRT加速：

trtexec --onnx=pruned_resnet.onnx \
        --saveEngine=resnet.plan \
        --workspace=1024

2. INT8量化：

model = torch.quantization.quantize_dynamic(
    model, {torch.nn.Linear}, dtype=torch.qint8
)

3. 多线程推理：

torch.set_num_threads(4)

在实际项目中，结合Network Slimming和后续优化，我们成功将ResNet-34的推理速度从23ms提升到12ms，满足了实时视频分析的需求。