模型剪枝及yolov5剪枝实践

文章目录

1、模型剪枝
1、稀疏化训练
2、模型剪枝
- 2.1 非结构化剪枝
- 2.2 结构化剪枝
- 2.3 一些疑惑：
- - 2.3.1 剪枝后参数量不变？
3、微调

【结构化剪枝掉点太多，不如一开始就选个小模型训练。非结构化剪枝只是checkpoint文件变小了，推理速度没有增益，精度还略微下降。总之目前没有认为模型剪枝在推理速度上有什么增益，比较鸡肋】

1、模型剪枝

模型剪枝的定义就不再赘述，剪枝就是将网络结构中对性能影响不大且权重值较小的节点去掉，从而降低模型的计算量，提高运算速率。

模型剪枝的流程：

正常训练：初始化模型并进行正常训练，获得一个初始模型。
稀疏化训练：在训练过程中引入稀疏化方法，使部分参数趋于零。
剪枝操作：根据稀疏化结果，去除零值或接近零的参数。
微调（Fine-tuning）：对剪枝后的模型进行重新训练，以恢复性能。

很好理解，在剪枝时哪些节点需要被裁剪，需要依靠一个指标。稀疏化就是在训练过程中加以约束，使得某些不重要的节点的参数趋于零，为后续剪枝提供依据。

虽然某些阶段的参数趋于零，但是仍然具有一定的作用，当裁剪后可能会对模型的性能产生较大的波动。所以再通过微调重新训练，恢复模型的性能。

1、稀疏化训练

数学层面的稀疏化：稀疏化的核心是通过某种方式（如正则化或剪枝）将权重矩阵中的绝大部分元素置零，使得非零元素的比例非常低。引申到模型中，就是将参数趋于零的过程。

实现稀疏化训练比较简单的就是引入L1正则化，所谓的L1正则化就是‌指向量中各个元素的绝对值之和。在模型训练过程中，引入L1正则化损失的稀疏化训练损失为：
$L=L_{ori}+λ||w||_1$
其实就是一个正常的损失函数 $L_{ori}$ （在分类中可能是交叉熵，在检测中可能是obj和bbox损失），加上一个特征值参数的绝对和。

可以想象到，模型越稀疏，权重参数的绝对和越小，其损失函数越小。也就是损失函数会驱使模型向着参数和小，稀疏程度高的方向更新。

2、模型剪枝

剪枝分为非结构化和结构化剪枝。

2.1 非结构化剪枝

非结构化就是将某些权重值置0，但是整个模型结构不改变。据说这种只是会减小checkpoint的储存大小，对模型速度没有太大的优化【未测试】，因为原先需要计算300步，现在仍然需要计算300步。结构化测试就是删除模型结构中的某些组件，比如删除某个卷积层之类的，这种可能就只需要计算200步，可以进行加速。

pytorch提供了很多工具方便剪枝，比如对卷积层权重矩阵中某些权重值裁剪：

def prune(model, amount=0.3):
    # Prune model to requested global sparsity
    import torch.nn.utils.prune as prune
    for name, m in model.named_modules():
        if isinstance(m, nn.Conv2d):
            prune.l1_unstructured(m, name='weight', amount=amount)  # 根据l1对m的weight进行裁剪，生成mask
            print(m.weight_mask[0])
            print(m.weight[0])
            prune.remove(m, 'weight')  # 根据mask对weight进行固化，并删除mask变量
    LOGGER.info(f'Model pruned to {sparsity(model):.3g} global sparsity')

这种方式会对卷积层的weight进行l1统计，从小到大排序将前30%的值置零。效果如下，这是模型中第一个卷积的权重，经过剪枝后，权重矩阵中有很多都被置0了，模型的大小也降低了一半。
在这里插入图片描述

2.2 结构化剪枝

非结构化只是对权重矩阵中某些个值置0，结构化则是可以将某些通道全部置零：

def structured_prune(model, amount=0.3):
    for name, m in model.named_modules():
        if isinstance(m, nn.Conv2d):
            prune.ln_structured(m, name='weight', amount=amount, n=2, dim=0)  # 对输出通道剪枝
            prune.remove(m, 'weight')  # 使剪枝永久化

关于ln_structured函数的解释：ln_structured 是 PyTorch 中 torch.nn.utils.prune 模块中的一个剪枝方法，它用于进行结构化剪枝（structured pruning）。该方法根据某种准则（例如 L2 范数、L1 范数）剪除模型中整个卷积核、通道、或其他结构化的单元，而不是单独的权重。结构化剪枝通常有助于减少模型的计算量和存储量，尤其是对于卷积神经网络（CNN）这样的模型。

经过结构化剪枝后，对(1,16,6,6)的卷积核，统计16个通道权重值，可以看到有16*0.3=4.8≈5个通道置0了。相对于非结构化剪枝，结构化剪枝在更高层面进行更大范围的剪枝。
在这里插入图片描述

2.3 一些疑惑：

2.3.1 剪枝后参数量不变？

如果使用total_params = sum(p.numel() for p in model.parameters())统计模型的参数量，会发现无论是非结构化剪枝还是结构化剪枝的数量和原始模型相比都没有下降。原因在于：1）非结构化中只是将某些单个值置0，整个模型结构不变；2）结构化剪枝将通道全部置0了，理论上可以去掉这个通道，比如上方的(1,16,6,6)卷积核可以去掉5个0值通道后，变成(1,11,6,6)。但是上方使用的pytorch提供的prune没有自动删除这些，只是用0值标记这个通道没用。所以可以用手动自行判断是否为0值通道然后删除，比如：

def structured_prune2(model, amount=0.3):
    for name, m in model.named_modules():
        if isinstance(m, nn.Conv2d):
            prune.ln_structured(m, name='weight', amount=amount, n=2, dim=0)  # 对输出通道剪枝
            prune.remove(m, 'weight')  # 使剪枝永久化

            weight = m.weight.data
            nonzero_channels = weight.abs().sum(dim=(1, 2, 3)) > 0  # 找出哪些通道仍然有效
            m.weight.data = weight[nonzero_channels]  # 只保留有效通道的权重

            if m.bias is not None:
                m.bias.data = m.bias.data[nonzero_channels]  # 更新偏置

可以看到原先(1,16,6,6)的卷积变成了(1,11,6,6)，此时再看参数量从1764118降到了1240371，下降了30%
在这里插入图片描述

3、微调

经过剪枝，模型的参数多多少少都发生了变化，肯定对模型的精度是有所影响的:

# 剪枝前
P        R        mAP50    mAP50-95
0.987    0.956    0.975    0.641

# 非结构化剪枝30%后
P        R        mAP50    mAP50-95
0.989    0.935    0.965    0.617

# 结构化剪枝30%后
P        R        mAP50    mAP50-95
0        0        0        0