模型剪枝实战|基于torch-pruning库代码对yolov8进行剪枝

torch-pruning库是一个开源的模型剪枝库，yolov8是是一个2年前较为先进的目标检测模型。在torch-pruning库中有很多模型剪枝案例，本文以yolov8剪枝代码为案例进行分析，代码路径在torch-pruning项目下examples\yolov8\yolov8_pruning.py。 本博文基于官方代码对coco128数据进行剪枝尝试，发现剪枝后的map有6个点的下降，这主要是coco128数据不够，同时官方的剪枝代码训练参数不够灵活。最终提出了修改意见，也对代码中关键部分进行分析。同时在2.4章节中进行剪枝测试，表明在模型剪枝中预训练权重还是有价值的。

torch-pruning项目介绍：https://blog.csdn.net/a486259/article/details/140421837
yolov8项目介绍：https://blog.csdn.net/a486259/article/details/135426696

1、基础准备

这里要求安装好了torch-cuda运行环境

1.1 torch-pruning库下载安装

打开https://github.com/VainF/Torch-Pruning，下载项目

然后在终端中，进入项目目录，并执行pip install -r requirements.txt 安装项目依赖库

然后在执行 pip install -e . ，将项目安装在当前目录下，并设置为editing模式。

验证安装：执行命令python -c “import torch_pruning”, 如果没有输出报错信息则表示安装成功。

1.2 ultralytics项目安装

torch-pruning库提供的yolov8剪枝代码只支持8.0.90之前的yolov8模型剪枝。因此，需要根据特定指令安装yolov8版本。

git clone https://github.com/ultralytics/ultralytics.git 
cp yolov8_pruning.py ultralytics/
cd ultralytics 
git checkout 44c7c3514d87a5e05cfb14dba5a3eeb6eb860e70 # for compatibility

也可以打开 https://github.com/ultralytics/ultralytics/tags?after=v8.0.91 下载早期版本

下载项目后执行 pip install -e .

1.3 代码拷贝

将torch-pruning项目下的 examples\yolov8\yolov8_pruning.py
拷贝到ultralytics项目根目录下。

2、运行效果

2.1 基础修改

若电脑配置较高（内存64g以上，显存6g以上）可以不用修改yolov8_pruning.py。
修改一：将第388行。model修改为yolov8s.pt，原先是yolov8m.pt

 parser = argparse.ArgumentParser()
    parser.add_argument('--model', default='yolov8s.pt', help='Pretrained pruning target model file')
    

修改二：在第285行后插入pruning_cfg[‘workers’] = 0，并将epochs修改为30，具体如下所示

2.2 剪枝效果

运行代码后，经过16次剪枝部署，将模型的channel剪枝为原来的0.5倍。

原始map信息，模型在coco数据集上训练，以下精度信息是在coco128上验证

剪枝后的模型都在coco128上训练与验证。
每次剪枝的flop与map变化如下：
第一次剪枝，剪枝后运行速度提升为1.07倍，map降低到0，重新训练30个epoch后map恢复到73



第五次剪枝，剪枝后运行速度提升为1.45倍，map降低到0，重新训练30个epoch后map恢复到74

第十次剪枝，剪枝后运行速度提升为2.03倍，map降低到20，重新训练30个epoch后map恢复到72


第十六次剪枝，剪枝后运行速度提升为2.95倍，map降低到0，重新训练30个epoch后map恢复到65


每一次的剪枝map与mac变化详情

2.3 剪枝特性分析

这里以第一次剪枝时的map精度为剪枝前精度，一开始map为73，最后一次剪枝后map为65。这存在显著的map下降，这并不是剪枝方法的问题。因为，训练与验证都是coco128，数据量比较少。如要使用该代码训练自己的剪枝模型，修改一下数据集配置接口。

第1次训练时的loss曲线如下所示，可以看到在第30个epoch时，map还是在平缓上升的，这表明30个epoch的训练有所不足的。

第16次训练时的loss曲线如下所示，可以看到在第30个epoch时，map还是在上升的，这表明30个epoch训练是严重不足的。同时loss曲线在第10个epoch后发生震荡，这表明此时学习率设置过大。

模型剪枝后后参数量变小，模型结构的变化对训练超参数的是有有所影响的，当前代码下基于同一训练策略得到的最终模型是不佳，需要进行二次训练。在此提出2点改进建议：
1、根据loss中后期震荡，表明要使用分阶段学习率
2、根据参数量变化，训练模型的batchsize应该适当调大

具体改进带代码如下所示，在最后一次剪枝时调大了batch与epoch，由于项目没有提供分阶段学习率，故而设置使用余弦学习率。关于各种学习率调度器，可以参考 https://blog.csdn.net/a486259/article/details/123074464

2.4 一次剪枝尝试

官方的剪枝代码通过16次剪枝尝试，将模型通道剪枝为原来的0.5倍。博主进行一次剪枝到0.5倍尝试，并将训练epoch扩展为原来的30倍（保持与16次迭代相同的训练资源）。这里是基于2.3中修改后的代码进行训练的。
可以发现这样子剪枝后，模型精度基本崩盘，精度上涨缓慢。

在第150个epoch时，map才30

在第200个epoch时，map才到44

到400多个epoch时，map达到了16次剪枝的水平。

直接对未训练权重进行剪枝，可以发现400多个epoch了，map还是很低。这表明剪枝还是有作用的。

3、关键代码分析

3.1 分次剪枝代码

首先 通过参数设置了iterative-steps=16，最终的目录剪枝率为0.5

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', default='yolov8s.pt', help='Pretrained pruning target model file')
    parser.add_argument('--cfg', default='default.yaml',
                        help='Pruning config file.'
                             ' This file should have same format with ultralytics/yolo/cfg/default.yaml')
    parser.add_argument('--iterative-steps', default=16, type=int, help='Total pruning iteration step')
    parser.add_argument('--target-prune-rate', default=0.5, type=float, help='Target pruning rate')
    parser.add_argument('--max-map-drop', default=0.2, type=float, help='Allowed maximum map drop after fine-tuning')

这里的分成剪枝代码并没有基于tp.pruner.GroupNormPruner中的iterative_steps参数进行设置，而是通过计算每个i步骤的剪枝率目标单独调用tp.pruner.GroupNormPruner进行剪枝设置。

3.2 替换或新增类成员方法

在yolov8_pruning.py中对多个yolov8模型类成员方法进行替换或新增，这里以save_model_v2函数为例。
首先在定义函数时，添加了一个self参数，并指定了参数类型

在替换原模型方法时，使用以下代码

3.3 替换模型结构

在原始的yolov8模型中有一个c2f层，由于使用了split操作,将一个模型的输出分割成2份给不同模块使用，导致不能正常的进行剪枝。

class C2f(nn.Module):
    """CSP Bottleneck with 2 convolutions."""

    def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
        super().__init__()
        self.c = int(c2 * e)  # hidden channels
        self.cv1 = Conv(c1, 2 * self.c, 1, 1)
        self.cv2 = Conv((2 + n) * self.c, c2, 1)  # optional act=FReLU(c2)
        self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))

    def forward(self, x):
        """Forward pass of a YOLOv5 CSPDarknet backbone layer."""
        y = list(self.cv1(x).chunk(2, 1))
        y.extend(m(y[-1]) for m in self.m)
        return self.cv2(torch.cat(y, 1))

    def forward_split(self, x):
        """Applies spatial attention to module's input."""
        y = list(self.cv1(x).split((self.c, self.c), 1))
        y.extend(m(y[-1]) for m in self.m)
        return self.cv2(torch.cat(y, 1))

通过replace_c2f_with_c2f_v2函数将c2f模块替换为c2f_v2模块。C2f_v2在定义上就将原先的cv1模块分解为cv0与cv1模块，然后调用transfer_weights函数进行权重转移。

def infer_shortcut(bottleneck):
    c1 = bottleneck.cv1.conv.in_channels
    c2 = bottleneck.cv2.conv.out_channels
    return c1 == c2 and hasattr(bottleneck, 'add') and bottleneck.add


class C2f_v2(nn.Module):
    # CSP Bottleneck with 2 convolutions
    def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
        super().__init__()
        self.c = int(c2 * e)  # hidden channels
        self.cv0 = Conv(c1, self.c, 1, 1)
        self.cv1 = Conv(c1, self.c, 1, 1)
        self.cv2 = Conv((2 + n) * self.c, c2, 1)  # optional act=FReLU(c2)
        self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))

    def forward(self, x):
        # y = list(self.cv1(x).chunk(2, 1))
        y = [self.cv0(x), self.cv1(x)]
        y.extend(m(y[-1]) for m in self.m)
        return self.cv2(torch.cat(y, 1))


def transfer_weights(c2f, c2f_v2):
    c2f_v2.cv2 = c2f.cv2
    c2f_v2.m = c2f.m

    state_dict = c2f.state_dict()
    state_dict_v2 = c2f_v2.state_dict()

    # Transfer cv1 weights from C2f to cv0 and cv1 in C2f_v2
    old_weight = state_dict['cv1.conv.weight']
    half_channels = old_weight.shape[0] // 2
    state_dict_v2['cv0.conv.weight'] = old_weight[:half_channels]
    state_dict_v2['cv1.conv.weight'] = old_weight[half_channels:]

    # Transfer cv1 batchnorm weights and buffers from C2f to cv0 and cv1 in C2f_v2
    for bn_key in ['weight', 'bias', 'running_mean', 'running_var']:
        old_bn = state_dict[f'cv1.bn.{bn_key}']
        state_dict_v2[f'cv0.bn.{bn_key}'] = old_bn[:half_channels]
        state_dict_v2[f'cv1.bn.{bn_key}'] = old_bn[half_channels:]

    # Transfer remaining weights and buffers
    for key in state_dict:
        if not key.startswith('cv1.'):
            state_dict_v2[key] = state_dict[key]

    # Transfer all non-method attributes
    for attr_name in dir(c2f):
        attr_value = getattr(c2f, attr_name)
        if not callable(attr_value) and '_' not in attr_name:
            setattr(c2f_v2, attr_name, attr_value)

    c2f_v2.load_state_dict(state_dict_v2)


def replace_c2f_with_c2f_v2(module):
    for name, child_module in module.named_children():
        if isinstance(child_module, C2f):
            # Replace C2f with C2f_v2 while preserving its parameters
            shortcut = infer_shortcut(child_module.m[0])
            c2f_v2 = C2f_v2(child_module.cv1.conv.in_channels, child_module.cv2.conv.out_channels,
                            n=len(child_module.m), shortcut=shortcut,
                            g=child_module.m[0].cv2.conv.groups,
                            e=child_module.c / child_module.cv2.conv.out_channels)
            transfer_weights(child_module, c2f_v2)
            setattr(module, name, c2f_v2)
        else:
            replace_c2f_with_c2f_v2(child_module)