基础

概念理解

白话：去掉模型里面没用的东西（权重参数、神经节点、通道、卷积层等）。

书面：剪枝是指在深度学习模型中，通过一定的评估标准，识别并移除对模型性能贡献较小的权重、神经元、通道或层等结构，从而减小模型的规模和复杂度，提高模型的运行效率和降低存储需求，同时也可能有助于减轻过拟合现象。

原理

白话：去掉对输出影响小的权重参数

书面语：冗余性假设、参数重要性评估（绝对值较小的权重、梯度较小的权重或对输出敏感度低的参数对模型的贡献较小）。

分类

• 按剪枝粒度划分 ：

  • 非结构化剪枝 ：权重剪枝，对神经网络中单个权重评估、剪枝，不考虑权重在神经网络中的位置和结构，只关注权重本身的数值大小或重要性。例如，将权重矩阵中绝对值最小的 10% 的权重置为零。这种方法可以更精细地控制模型的剪枝程度，但在模型进行推断时，由于需要处理稀疏矩阵，因此对硬件和软件的支持要求较高，否则可能难以获得明显的速度提升。

  • 结构化剪枝 ：移除模型中整个结构化的单元，如卷积神经网络中的滤波器、通道或整个层等。比如直接删除卷积层的某些通道或全连接层的某些神经元。这种剪枝方式能够使模型结构更加紧凑和规则，便于在各种硬件平台上实现加速，但其剪枝的粒度相对较粗，可能会对模型的性能造成较大影响。不过，一旦剪枝完成，模型的推理速度将得到显著提高，因为在计算过程中可以完全跳过被剪除的结构。

• 按剪枝过程划分 ：

  • 一次性剪枝 ：在模型训练完成后，根据设定的剪枝规则和比例，一次性地将不重要的参数剪除。这种方法简单直接，但可能会导致模型性能的较大波动，尤其是在剪枝比例较高时，需要后续充分的微调来恢复模型的性能。

  • 迭代式剪枝 ：先对模型进行少量的剪枝，然后对剪枝后的模型进行微调，使模型在剪枝后的结构下重新学习和调整参数；接着再进行下一轮的剪枝和微调，如此循环往复，逐步提高模型的稀疏率或压缩程度。与一次性剪枝相比，迭代式剪枝能够更好地平衡剪枝程度和模型性能之间的关系，通常可以获得更高的压缩率，同时保持较好的模型性能。

• 按参数选择方式划分 ：

  • 基于权重大小的剪枝 ：认为绝对值较小的权重对模型输出的贡献较小，可优先剪除。通常会设定一个全局或分层的剪枝率，然后将模型中绝对值最小的相应比例的权重置为零。这种方法实现简单，应用广泛，但可能会忽略一些在特定情况下虽然权重值小但对模型泛化能力有重要影响的连接。

  • 基于梯度的剪枝 ：认为梯度值较小的权重对损失函数的影响较小，对模型的输出不那么重要，因此可以剪除这些权重。在训练过程中，通过权重的梯度信息来评估每个权重的重要性，将梯度绝对值较小的权重剪枝。与基于权重大小的剪枝相比，该方法能够更直接地考虑权重在模型训练过程中的实际作用，但需要额外计算梯度信息，计算成本相对较高。

  • 基于敏感度的剪枝 ：通过计算每个权重或神经元对模型输出的敏感度来确定其重要性。敏感度越高，说明该权重或神经元对模型输出的影响越大，应保留；反之则可剪除。这种方法需要对模型进行大量的敏感度分析计算，计算成本较高，但在确定参数重要性方面更加准确。

  • 基于聚类中心的剪枝 ：对网络中的权重进行聚类分析，将相似的权重归为一类，并用该类的中心值来代表这一类权重。然后，基于某种准则，如权重与聚类中心的距离等，确定哪些权重可以被剪除。这种方法可以去除冗余的相似权重，降低模型的复杂度，但聚类过程本身可能会导致模型性能的一定下降。

  • 基于 L1/L2 范数的剪枝 ：除了直接基于数值大小，还可以利用权重的范数来衡量其重要性。L1 范数倾向于产生更加稀疏的解，而 L2 范数则更倾向于减小权重的整体幅度。在非结构化剪枝中，通常是对单个权重计算其 L1 或 L2 范数（对于单个权重，两者本质上都是绝对值），根据剪枝率，移除范数最小的权重。

  • 基于信息熵的剪枝 ：计算每个权重或特征的重要性信息熵，信息熵较低的权重对输出不敏感，可以被移除。

  • 基于低秩分解的剪枝 ：通过对权重矩阵进行低秩分解，将原权重矩阵近似表示为几个低秩矩阵的乘积，从而降低权重矩阵的维度，实现模型的压缩。这种方法可以在一定程度上减少模型的参数量和计算量，但可能会改变模型的结构和参数分布，需要进一步的微调和优化。

  • 基于重构误差的剪枝 ：在剪枝过程中，以最小化模型的重构误差为目标函数，通过不断调整剪枝策略，寻找在满足一定压缩率要求下的最优剪枝方案，使剪枝后的模型在保持原模型结构和功能的基础上，尽可能减少性能损失。

  • 基于平均失活率的剪枝 ：对于卷积神经网络中的通道剪枝，平均失活率是一种常用的评估指标。在剪枝过程中，计算每个通道在训练数据上被激活的平均值，若平均失活率较高，则说明该通道对于特征提取的作用较小，可以被剪除。这种方法能够较好地适应不同的卷积神经网络结构，但需要设计不同的剪枝策略来适应不同的网络结构。

  • 基于泰勒展开式的剪枝 ：利用泰勒展开式对模型的损失函数进行近似，计算每个权重对损失函数的贡献程度，从而确定权重的重要性。这种方法能够更准确地评估权重对模型性能的影响，但计算成本较高，且需要对模型的数学结构有较深入的理解。

  • 基于强化学习的剪枝 ：将剪枝过程视为一个强化学习问题，定义一个搜索空间，包含所有可能的剪枝决策（如每个层可以剪除的滤波器或通道的比例等）。然后训练一个强化学习 Agent，该 Agent 通过与环境（即待剪枝的模型）交互来学习最优的剪枝策略。Agent 的状态通常包括当前模型的性能指标和结构信息，而 Reward 函数则结合了剪枝后模型的性能和模型大小等因素，目标是找到一个在两者之间取得最佳平衡的剪枝策略。

  • 基于神经架构搜索的剪枝 ：在使用神经架构搜索进行剪枝时，通常会将剪枝率或要剪除的结构类型视为搜索空间的一部分。搜索算法（如进化算法、基于梯度的搜索等）会探索不同的剪枝配置，并在验证集上评估它们的性能（通常是精度和模型大小的权衡）。通过迭代搜索，神经架构搜索可以找到在满足特定资源约束下性能最优的剪枝模型结构。

组合剪枝

可以挨个使用上述方法对一个模型进行多次剪枝。组合使用多种剪枝方法能够在不同层面上优化模型，从而获得更优的性能和效率。以下是关于组合使用剪枝方法的一些细节：

组合剪枝的可行性

多种剪枝方法关注模型中的不同冗余信息，如非结构化剪枝关注权重本身，结构化剪枝关注模型的结构。组合使用这些方法可以从多个角度优化模型，理论上是可行的。

组合剪枝的注意事项

• 剪枝顺序：剪枝方法的先后顺序可能会影响最终效果。例如，如果先进行结构化剪枝，改变了模型的结构，可能会影响后续非结构化剪枝的选择。因此，需要根据模型的具体情况和目标来确定剪枝顺序。

• 模型的完整性和功能：在组合剪枝过程中，要确保模型的完整性和功能不受影响。过度剪枝可能导致模型性能下降或功能异常，因此需要在剪枝过程中不断评估模型的性能。

• 剪枝率的控制：每次剪枝时，都应合理控制剪枝率。剪枝率过高可能导致模型性能急剧下降，因此需要逐步进行剪枝，并在每次剪枝后对模型进行微调和评估。

实际操作建议

• 分步剪枝：可以先应用一种剪枝方法，如非结构化剪枝，对模型进行初步优化，然后再应用结构化剪枝或其他方法进行进一步优化。

• 微调和评估：在每次剪枝后，都要对模型进行微调，以恢复或提升模型的性能，并通过评估确定是否达到预期效果。如果效果不理想，可能需要调整剪枝策略或回退到之前的模型版本。

• 记录和分析：记录每次剪枝的操作和结果，分析不同剪枝方法对模型性能的影响，以便优化剪枝策略。

具体实现方式

对模型进行剪枝操作，并不是直接修改 yolov8.yaml 文件来实现的。yolov8.yaml 文件主要是用于定义 YOLOv8 模型的结构，包括各个层的类型、参数等，只是一个模型结构的定义文件，不包含模型的权重信息。剪枝操作需要对模型的权重或结构进行实际的修改，这些操作是在代码中通过深度学习框架（如 PyTorch）实现的，而不是通过修改配置文件来完成的。

剪枝操作是在训练或推理代码中进行的：

1. 在训练过程中进行剪枝：

   • 需要修改训练代码，在训练过程中应用剪枝算法（如结构化剪枝、非结构化剪枝等），动态地调整模型的参数或结构调整。

2. 在推理过程中进行剪枝：

   • 对已训练好的模型进行剪枝操作，通常需要加载模型，然后应用剪枝算法对模型进行处理，最后保存剪枝后的模型。

在训练代码中集成剪枝算法的步骤示例：

选择剪枝库

• 可选择torch-pruning等剪枝库，其提供了丰富的剪枝算法和工具，方便对模型进行剪枝操作。

定义剪枝策略

• 剪枝比例：设置剪枝比例，如将模型的参数量减少到原来的一定比例（如减少到80%），以减小模型规模，提升推理速度。

• 迭代步骤：确定剪枝的迭代步骤，如分4次迭代完成剪枝，逐步减少模型的参数量。

• 目标剪枝率：设置目标剪枝率，如剪枝率为20%，即剪除20%的参数或通道。

修改训练代码

• 加载模型：加载需要剪枝的YOLOv8模型。

• 替换模块：将模型中的C2f模块替换为自定义的C2f_v2模块，以适应剪枝操作。

• 初始化权重：对替换后的模型进行权重初始化。

• 定义剪枝器：根据选择的剪枝算法，定义剪枝器。

• 剪枝迭代：在每次迭代中，对模型进行剪枝操作，然后进行微调。

• 保存模型：保存剪枝后的模型，用于后续的推理任务。

示例代码详解

1. 导入必要的库和模块

import argparse
import math
import os
from copy import deepcopy
from datetime import datetime
from pathlib import Path
from typing import List, Union

import numpy as np
import torch
import torch.nn as nn
from matplotlib import pyplot as plt
from ultralytics import YOLO, __version__
from ultralytics.nn.modules import Detect, C2f, Conv, Bottleneck
from ultralytics.nn.tasks import attempt_load_one_weight
from ultralytics.yolo.engine.model import TASK_MAP
from ultralytics.yolo.engine.trainer import BaseTrainer
from ultralytics.yolo.utils import yaml_load, LOGGER, RANK, DEFAULT_CFG_DICT, DEFAULT_CFG_KEYS
from ultralytics.yolo.utils.checks import check_yaml
from ultralytics.yolo.utils.torch_utils import initialize_weights, de_parallel
import torch_pruning as tp

• argparse：用于解析命令行参数。

• math：用于数学计算。

• os：用于操作系统相关操作，如设置环境变量。

• copy：用于深拷贝对象。

• datetime：用于处理日期和时间。

• pathlib：用于路径操作。

• numpy：用于数值计算。

• torch：PyTorch深度学习框架。

• matplotlib.pyplot：用于绘图。

• ultralytics：YOLOv8模型相关模块。

• torch_pruning：用于模型剪枝的库。

2. 定义保存剪枝性能图的函数

def save_pruning_performance_graph(x, y1, y2, y3):
    # 保存剪枝性能图
    try:
        plt.style.use("ggplot")
    except:
        pass
    x, y1, y2, y3 = np.array(x), np.array(y1), np.array(y2), np.array(y3)
    y2_ratio = y2 / y2[0]
    fig, ax = plt.subplots(figsize=(8, 6))
    ax.set_xlabel('Pruning Ratio')
    ax.set_ylabel('mAP')
    ax.plot(x, y1, label='recovered mAP')
    ax.scatter(x, y1)
    ax.plot(x, y3, color='tab:gray', label='pruned mAP')
    ax.scatter(x, y3, color='tab:gray')
    ax2 = ax.twinx()
    ax2.set_ylabel('MACs')
    ax2.plot(x, y2_ratio, color='tab:orange', label='MACs')
    ax2.scatter(x, y2_ratio, color='tab:orange')
    lines, labels = ax.get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()
    ax2.legend(lines + lines2, labels + labels2, loc='best')
    ax.set_xlim(105, -5)
    ax.set_ylim(0, max(y1) + 0.05)
    ax2.set_ylim(0.05, 1.05)
    max_y1_idx = np.argmax(y1)
    min_y1_idx = np.argmin(y1)
    max_y2_idx = np.argmax(y2)
    min_y2_idx = np.argmin(y2)
    max_y1 = y1[max_y1_idx]
    min_y1 = y1[min_y1_idx]
    max_y2 = y2_ratio[max_y2_idx]
    min_y2 = y2_ratio[min_y2_idx]
    ax.text(x[max_y1_idx], max_y1 - 0.05, f'max mAP = {max_y1:.2f}', fontsize=10)
    ax.text(x[min_y1_idx], min_y1 + 0.02, f'min mAP = {min_y1:.2f}', fontsize=10)
    ax2.text(x[max_y2_idx], max_y2 - 0.05, f'max MACs = {max_y2 * y2[0] / 1e9:.2f}G', fontsize=10)
    ax2.text(x[min_y2_idx], min_y2 + 0.02, f'min MACs = {min_y2 * y2[0] / 1e9:.2f}G', fontsize=10)
    plt.title('Comparison of mAP and MACs with Pruning Ratio')
    plt.savefig('pruning_perf_change.png')

• 功能：保存剪枝过程中的性能变化图，包括mAP（平均精度均值）和MACs（乘加运算次数）。

• 参数：

  • x：剪枝比例列表。

  • y1：恢复后的mAP值列表。

  • y2：MACs值列表。

  • y3：剪枝后的mAP值列表。

• 实现：

  • 使用matplotlib绘制双Y轴图表，一个轴表示mAP，另一个轴表示MACs。

  • 标注最大和最小值，并保存图表。

3. 判断是否可以直接连接瓶颈层的输入和输出

瓶颈层（Bottleneck）是深度学习中用于优化模型结构的一种设计模式，常见于卷积神经网络（CNN）。它通过减少数据在层间的传输量，降低计算复杂度并提升模型效率。

工作原理

• 通道数压缩：第一个1×1卷积减少输入通道数，降低后续3×3卷积的计算量。

• 特征提取：3×3卷积在较低通道数下提取特征，减少计算量的同时保留关键信息。

• 通道数恢复：第三个1×1卷积恢复通道数至与输入相同，便于与输入直接相加。

假设你正在装修房子，想把两个房间直接连通。瓶颈层的输入和输出就像这两个房间的门。你需要检查两个房间的门的大小（通道数c1和c2）是否一致，并且检查房间里是否有允许直接连通的条件（“add”属性）。如果门的大小一致且允许直接连通，那你就可以直接打通这两个房间；否则，你就不能直接连通，而需要经过房间内的其他结构。

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

• 逐句解释

• def infer_shortcut(bottleneck):

  • 定义一个名为infer_shortcut的函数，输入参数是bottleneck，即瓶颈层。

• c1 = bottleneck.cv1.conv.in_channels

  • 获取瓶颈层中第一个卷积层（cv1）的输入通道数，并将其赋值给变量c1。

• c2 = bottleneck.cv2.conv.out_channels

  • 获取瓶颈层中第二个卷积层（cv2）的输出通道数，并将其赋值给变量c2。

• return c1 == c2 and hasattr(bottleneck, 'add') and bottleneck.add

  • 判断c1是否等于c2：即瓶颈层输入和输出的通道数是否一致。

  • 判断瓶颈层是否具有“add”属性：hasattr(bottleneck, 'add')。

  • 判断瓶颈层的“add”属性是否为真：bottleneck.add。

  • 只有当以上三个条件都满足时，函数才返回True，表示可以直连；否则返回False。

4. 定义自定义的C2f模块

自定义的C2f_v2模块继承了原始C2f模块的功能，用于替换YOLOv8模型中的原始C2f模块，同时对结构进行了调整，以适应剪枝操作。

class C2f_v2(nn.Module):
    # CSP Bottleneck with 2 convolutions
    def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):
        super().__init__()
        self.c = int(c2 * e)
        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)
        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 = [self.cv0(x), self.cv1(x)]
        y.extend(m(y[-1]) for m in self.m)
        return self.cv2(torch.cat(y, 1))

• 逐句解释

  • class C2f_v2(nn.Module):

    • 定义一个名为C2f_v2的类，继承自nn.Module，表示这是一个PyTorch神经网络模块。

  • def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):

    • 定义构造函数，初始化C2f_v2模块。

    • c1：输入通道数。

    • c2：输出通道数。

    • n：Bottleneck模块的数量，默认为1。

    • shortcut：是否使用快捷连接，默认为False。

    • g：分组数，默认为1。

    • e：扩展系数，默认为0.5。

  • super().__init__()

    • 调用父类nn.Module的构造函数，完成初始化。

  • self.c = int(c2 * e)

    • 计算中间层的通道数self.c，等于输出通道数c2乘以扩展系数e并取整。

  • self.cv0 = Conv(c1, self.c, 1, 1)

    • 定义第一个卷积层cv0，将输入通道数c1转换为中间通道数self.c，卷积核大小为1×1，步长为1。

  • self.cv1 = Conv(c1, self.c, 1, 1)

    • 定义第二个卷积层cv1，同样将输入通道数c1转换为中间通道数self.c，卷积核大小为1×1，步长为1。

  • self.cv2 = Conv((2 + n) * self.c, c2, 1)

    • 定义第三个卷积层cv2，将多个特征图拼接后的通道数(2 + n) * self.c转换为输出通道数c2，卷积核大小为1×1。

  • self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))

    • 定义一个Bottleneck模块列表self.m，包含n个Bottleneck模块。每个Bottleneck模块的输入和输出通道数均为self.c，使用快捷连接shortcut，分组数为g，卷积核大小为3×3，扩展系数为1.0。

  • def forward(self, x):

    • 定义前向传播函数。

  • y = [self.cv0(x), self.cv1(x)]

    • 将输入x分别通过cv0和cv1卷积层，得到两个特征图，并将它们放入列表y中。

  • y.extend(m(y[-1]) for m in self.m)

    • 将输入依次通过每个Bottleneck模块m，并将输出追加到列表y中。

  • return self.cv2(torch.cat(y, 1))

    • 将列表y中的所有特征图在通道维度上进行拼接（torch.cat(y, 1)），然后通过cv2卷积层，得到最终输出。

• 实现：

  • 定义了三个卷积层（cv0, cv1, cv2）和一个Bottleneck模块列表（m）。

  • 前向传播时，将输入通过cv0和cv1，然后依次通过Bottleneck模块，最后将所有特征图拼接并通过cv2。

5. 转移权重函数

将原始C2f模块中的权重转移到自定义的C2f_v2模块中，确保在替换模块后，C2f_v2模块能够继承原始模块的权重，从而保持模型性能的连续性。

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()
    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:]
    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:]
    for key in state_dict:
        if not key.startswith('cv1.'):
            state_dict_v2[key] = state_dict[key]
    c2f_v2.load_state_dict(state_dict_v2)

原理

• 直接赋值：对于cv2和m模块，直接将原始模块的对应部分赋值给C2f_v2模块。

• 权重切分：对于cv1的权重，将其切分为两部分，分别赋值给cv0和cv1。

• 批量归一化参数切分：对于cv1的批量归一化（BatchNorm）参数，同样将其切分为两部分，分别赋值给cv0和cv1。

逐句解释

1. def transfer_weights(c2f, c2f_v2):

   • 定义一个名为transfer_weights的函数，输入参数是原始的C2f模块和自定义的C2f_v2模块。

2. c2f_v2.cv2 = c2f.cv2

   • 将原始C2f模块的cv2卷积层直接赋值给C2f_v2模块的cv2卷积层。

3. c2f_v2.m = c2f.m

   • 将原始C2f模块的Bottleneck模块列表m直接赋值给C2f_v2模块的m。

4. state_dict = c2f.state_dict()

   • 获取原始C2f模块的权重字典，存储在state_dict中。

5. state_dict_v2 = c2f_v2.state_dict()

   • 获取C2f_v2模块的权重字典，存储在state_dict_v2中。

6. old_weight = state_dict['cv1.conv.weight']

   • 获取原始C2f模块中cv1卷积层的权重，存储在old_weight中。

7. half_channels = old_weight.shape[0] // 2

   • 计算old_weight的通道数的一半，存储在half_channels中。

8. state_dict_v2['cv0.conv.weight'] = old_weight[:half_channels]

   • 将old_weight的前half_channels通道赋值给C2f_v2模块中cv0卷积层的权重。

9. state_dict_v2['cv1.conv.weight'] = old_weight[half_channels:]

   • 将old_weight的后half_channels通道赋值给C2f_v2模块中cv1卷积层的权重。

10. for bn_key in ['weight', 'bias', 'running_mean', 'running_var']:

    • 遍历批量归一化（BatchNorm）的参数类型。

11. old_bn = state_dict[f'cv1.bn.{bn_key}']

    • 获取原始C2f模块中cv1批量归一化层的对应参数，存储在old_bn中。

12. state_dict_v2[f'cv0.bn.{bn_key}'] = old_bn[:half_channels]

    • 将old_bn的前half_channels部分赋值给C2f_v2模块中cv0批量归一化层的对应参数。

13. state_dict_v2[f'cv1.bn.{bn_key}'] = old_bn[half_channels:]

    • 将old_bn的后half_channels部分赋值给C2f_v2模块中cv1批量归一化层的对应参数。

14. for key in state_dict:

    • 遍历原始C2f模块的权重字典中的所有键。

15. if not key.startswith('cv1.'):

    • 如果键不以cv1.开头，即不属于cv1的权重。

16. state_dict_v2[key] = state_dict[key]

    • 将该键对应的权重直接赋值给C2f_v2模块的权重字典。

17. c2f_v2.load_state_dict(state_dict_v2)

    • 将更新后的权重字典state_dict_v2加载到C2f_v2模块中。

6. 替换C2f模块为C2f_v2模块

递归地遍历模型的所有子模块，找到每个C2f模块；

将它们替换为新的C2f_v2模块；

调用transfer_weights函数，将原始C2f模块的权重转移到新的C2f_v2模块中。

def replace_c2f_with_c2f_v2(module):
    for name, child_module in module.named_children():
        if isinstance(child_module, C2f):
            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)

逐句解释

1. def replace_c2f_with_c2f_v2(module):

   • 定义一个函数replace_c2f_with_c2f_v2，输入参数是一个PyTorch神经网络模块module。

2. for name, child_module in module.named_children():

   • 遍历输入模块module的所有直接子模块。name是子模块的名称，child_module是子模块本身。

3. if isinstance(child_module, C2f):

   • 检查当前子模块是否是C2f类的实例。如果是，则进行替换操作。

4. shortcut = infer_shortcut(child_module.m[0])

   • 调用infer_shortcut函数判断是否可以直接连接瓶颈层的输入和输出。child_module.m[0]是第一个Bottleneck模块。

5. c2f_v2 = C2f_v2(...)

   • 创建一个新的C2f_v2模块实例，并初始化其参数：

     • child_module.cv1.conv.in_channels：输入通道数。

     • child_module.cv2.conv.out_channels：输出通道数。

     • n=len(child_module.m)：Bottleneck模块的数量。

     • shortcut=shortcut：是否使用快捷连接。

     • g=child_module.m[0].cv2.conv.groups：分组数。

     • e=child_module.c / child_module.cv2.conv.out_channels：扩展系数。

6. transfer_weights(child_module, c2f_v2)

   • 调用transfer_weights函数，将原始C2f模块的权重转移到新的C2f_v2模块中。

7. setattr(module, name, c2f_v2)

   • 使用setattr函数将新的C2f_v2模块替换到原始模块的位置。name是原始子模块的名称，c2f_v2是新的模块实例。

8. else:

   • 如果当前子模块不是C2f类的实例，则递归调用replace_c2f_with_c2f_v2函数，继续检查该子模块的子模块。

9. replace_c2f_with_c2f_v2(child_module)

   • 递归调用函数，处理当前子模块的子模块。

7. 定义剪枝函数

对模型进行剪枝操作。

原理

• 剪枝策略：通过多次迭代剪枝，逐步减少模型的参数量。每次迭代中，使用剪枝算法确定并移除不重要的通道或权重。

• 权重转移：在替换模块后，将原始模块的权重转移到新的模块中，确保模型性能的连续性。

• 微调：每次剪枝后，对模型进行微调，恢复模型性能。

def prune(args):
    base_name = 'prune/' + str(datetime.now()) + '/'
    model = YOLO(args.model)
    model.__setattr__("train_v2", train_v2.__get__(model))
    pruning_cfg = yaml_load(check_yaml(args.cfg))
    batch_size = pruning_cfg['batch']
    pruning_cfg['data'] = "./ultralytics/datasets/soccer.yaml"
    pruning_cfg['epochs'] = 4
    model.model.train()
    replace_c2f_with_c2f_v2(model.model)
    initialize_weights(model.model)
    for name, param in model.model.named_parameters():
        param.requires_grad = True
    example_inputs = torch.randn(1, 3, pruning_cfg["imgsz"], pruning_cfg["imgsz"]).to(model.device)
    macs_list, nparams_list, map_list, pruned_map_list = [], [], [], []
    base_macs, base_nparams = tp.utils.count_ops_and_params(model.model, example_inputs)
    pruning_cfg['name'] = base_name + f"baseline_val"
    pruning_cfg['batch'] = 128
    validation_model = deepcopy(model)
    metric = validation_model.val(**pruning_cfg)
    init_map = metric.box.map
    macs_list.append(base_macs)
    nparams_list.append(100)
    map_list.append(init_map)
    pruned_map_list.append(init_map)
    print(f"Before Pruning: MACs={base_macs / 1e9: .5f} G, #Params={base_nparams / 1e6: .5f} M, mAP={init_map: .5f}")
    ch_sparsity = 1 - math.pow((1 - args.target_prune_rate), 1 / args.iterative_steps)
    for i in range(args.iterative_steps):
        model.model.train()
        for name, param in model.model.named_parameters():
            param.requires_grad = True
        ignored_layers = []
        unwrapped_parameters = []
        for m in model.model.modules():
            if isinstance(m, (Detect,)):
                ignored_layers.append(m)
        example_inputs = example_inputs.to(model.device)
        pruner = tp.pruner.GroupNormPruner(model.model, example_inputs,
                                           importance=tp.importance.GroupNormImportance(),
                                           iterative_steps=1,
                                           ch_sparsity=ch_sparsity,
                                           ignored_layers=ignored_layers,
                                           unwrapped_parameters=unwrapped_parameters)
        pruner.step()
        pruning_cfg['name'] = base_name + f"step_{i}_pre_val"
        pruning_cfg['batch'] = 128
        validation_model.model = deepcopy(model.model)
        metric = validation_model.val(**pruning_cfg)
        pruned_map = metric.box.map
        pruned_macs, pruned_nparams = tp.utils.count_ops_and_params(pruner.model, example_inputs.to(model.device))
        current_speed_up = float(macs_list[0]) / pruned_macs
        print(f"After pruning iter {i + 1}: MACs={pruned_macs / 1e9} G, #Params={pruned_nparams / 1e6: .5f} M, mAP={pruned_map: .5f}")
        macs_list.append(pruned_macs)
        nparams_list.append(100 - (i + 1) * args.target_prune_rate * 100 / args.iterative_steps)
        map_list.append(init_map)
        pruned_map_list.append(pruned_map)
    save_pruning_performance_graph(nparams_list, map_list, macs_list, pruned_map_list)

1. def prune(args):

   • 定义一个名为prune的函数，输入参数是命令行参数args。

2. base_name = 'prune/' + str(datetime.now()) + '/'

   • 创建一个基于当前时间的目录名称，用于存储剪枝过程中的模型和结果。

3. model = YOLO(args.model)

   • 加载YOLO模型，模型路径由命令行参数args.model指定。

4. model.__setattr__("train_v2", train_v2.__get__(model))

   • 将自定义的训练方法train_v2绑定到模型对象上。

5. pruning_cfg = yaml_load(check_yaml(args.cfg))

   • 加载剪枝配置文件，配置文件路径由命令行参数args.cfg指定。

6. batch_size = pruning_cfg['batch']

   • 获取剪枝配置中的批量大小。

7. pruning_cfg['data'] = "./ultralytics/datasets/soccer.yaml"

   • 设置数据集配置文件路径。

8. pruning_cfg['epochs'] = 4

   • 设置剪枝过程中的训练轮数为4。

9. model.model.train()

   • 将模型设置为训练模式。

10. replace_c2f_with_c2f_v2(model.model)

    • 调用replace_c2f_with_c2f_v2函数，将模型中的C2f模块替换为C2f_v2模块。

11. initialize_weights(model.model)

    • 初始化模型的权重。

12. for name, param in model.model.named_parameters():

    • 遍历模型的所有参数。

13. param.requires_grad = True

    • 将所有参数的梯度计算设置为True，表示这些参数在训练过程中需要更新。

14. example_inputs = torch.randn(1, 3, pruning_cfg["imgsz"], pruning_cfg["imgsz"]).to(model.device)

    • 创建一个示例输入张量，用于计算模型的MACs（乘加运算次数）和参数量。

15. macs_list, nparams_list, map_list, pruned_map_list = [], [], [], []

    • 初始化列表，用于存储剪枝过程中的MACs、参数量和mAP（平均精度均值）值。

16. base_macs, base_nparams = tp.utils.count_ops_and_params(model.model, example_inputs)

    • 计算模型的初始MACs和参数量。

17. pruning_cfg['name'] = base_name + f"baseline_val"

    • 设置验证模型的名称。

18. pruning_cfg['batch'] = 128

    • 设置验证时的批量大小为128。

19. validation_model = deepcopy(model)

    • 创建模型的深拷贝，用于验证。

20. metric = validation_model.val(**pruning_cfg)

    • 使用验证模型进行验证，获取性能指标。

21. init_map = metric.box.map

    • 获取初始的mAP值。

22. macs_list.append(base_macs)

    • 将初始MACs值添加到列表中。

23. nparams_list.append(100)

    • 将初始参数量（以百分比表示）添加到列表中。

24. map_list.append(init_map)

    • 将初始mAP值添加到列表中。

25. pruned_map_list.append(init_map)

    • 将初始mAP值添加到剪枝后的mAP列表中。

26. print(f"Before Pruning: MACs={base_macs / 1e9: .5f} G, #Params={base_nparams / 1e6: .5f} M, mAP={init_map: .5f}")

    • 打印剪枝前的模型性能指标。

27. ch_sparsity = 1 - math.pow((1 - args.target_prune_rate), 1 / args.iterative_steps)

    • 计算每次迭代的通道稀疏度。

28. for i in range(args.iterative_steps):

    • 开始迭代剪枝过程，迭代次数由命令行参数args.iterative_steps指定。

29. model.model.train()

    • 将模型设置为训练模式。

30. for name, param in model.model.named_parameters():

    • 遍历模型的所有参数。

31. param.requires_grad = True

    • 将所有参数的梯度计算设置为True。

32. ignored_layers = []

    • 初始化一个列表，用于存储不进行剪枝的层。

33. unwrapped_parameters = []

    • 初始化一个列表，用于存储未包装的参数。

34. for m in model.model.modules():

    • 遍历模型的所有模块。

35. if isinstance(m, (Detect,)):

    • 检查模块是否是Detect模块。

36. ignored_layers.append(m)

    • 将Detect模块添加到ignored_layers列表中，表示这些层不进行剪枝。

37. example_inputs = example_inputs.to(model.device)

    • 将示例输入张量移动到模型所在的设备（CPU或GPU）。

38. pruner = tp.pruner.GroupNormPruner(model.model, example_inputs, ...)

    • 创建一个剪枝器实例，使用GroupNormPruner进行剪枝：

      • model.model：需要剪枝的模型。

      • example_inputs：示例输入张量。

      • importance=tp.importance.GroupNormImportance()：使用组归一化重要性评估方法。

      • iterative_steps=1：每次迭代的剪枝步数。

      • ch_sparsity=ch_sparsity：通道稀疏度。

      • ignored_layers=ignored_layers：不进行剪枝的层。

      • unwrapped_parameters=unwrapped_parameters：未包装的参数。

39. pruner.step()

    • 执行一次剪枝操作。

40. pruning_cfg['name'] = base_name + f"step_{i}_pre_val"

    • 设置当前迭代步骤的验证模型名称。

41. validation_model.model = deepcopy(model.model)

    • 深拷贝当前模型，用于验证。

42. metric = validation_model.val(**pruning_cfg)

    • 使用验证模型进行验证，获取性能指标。

43. pruned_map = metric.box.map

    • 获取剪枝后的mAP值。

44. pruned_macs, pruned_nparams = tp.utils.count_ops_and_params(pruner.model, example_inputs.to(model.device))

    • 计算剪枝后的模型的MACs和参数量。

45. current_speed_up = float(macs_list[0]) / pruned_macs

    • 计算当前剪枝步骤后的加速比。

46. print(f"After pruning iter {i + 1}: MACs={pruned_macs / 1e9} G, #Params={pruned_nparams / 1e6: .5f} M, mAP={pruned_map: .5f}")

    • 打印当前剪枝步骤后的模型性能指标。

47. macs_list.append(pruned_macs)

    • 将剪枝后的MACs值添加到列表中。

48. nparams_list.append(100 - (i + 1) * args.target_prune_rate * 100 / args.iterative_steps)

    • 将剪枝后的参数量（以百分比表示）添加到列表中。

49. map_list.append(init_map)

    • 将初始mAP值添加到列表中。

50. pruned_map_list.append(pruned_map)

    • 将剪枝后的mAP值添加到列表中。

51. save_pruning_performance_graph(nparams_list, map_list, macs_list, pruned_map_list)

    • 调用函数保存剪枝过程中的性能变化图。

8. 主函数

定义主函数，用于解析命令行参数并调用剪枝函数，启动剪枝操作。

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--cfg", type=str, default="yolov8.yaml", help="model.yaml path")
    parser.add_argument("--data", type=str, default="coco128.yaml", help="dataset.yaml path")
    parser.add_argument("--model", type=str, default="yolov8n.pt", help="initial weights path")
    parser.add_argument("--target-prune-rate", type=float, default=0.2, help="target prune rate")
    parser.add_argument("--iterative-steps", type=int, default=4, help="iterative steps for pruning")
    args = parser.parse_args()
    prune(args)

1. if __name__ == "__main__":

   • 检查是否直接运行当前脚本。如果是，则执行以下代码。

     • __name__ 是 Python 中的一个特殊变量，表示当前模块的名称。当脚本直接运行时，__name__ 被设置为 "__main__"。

2. parser = argparse.ArgumentParser()

   • 创建一个ArgumentParser对象，用于解析命令行参数。

     • ArgumentParser 是 argparse 库中的一个类，用于处理命令行参数。

3. parser.add_argument("--cfg", type=str, default="yolov8.yaml", help="model.yaml path")

   • 添加一个命令行参数--cfg，用于指定模型配置文件的路径。

     • type=str：参数类型为字符串。

     • default="yolov8.yaml"：默认值为"yolov8.yaml"。

     • help="model.yaml path"：帮助信息，描述该参数的用途。

4. parser.add_argument("--data", type=str, default="coco128.yaml", help="dataset.yaml path")

   • 添加一个命令行参数--data，用于指定数据集配置文件的路径。

     • type=str：参数类型为字符串。

     • default="coco128.yaml"：默认值为"coco128.yaml"。

     • help="dataset.yaml path"：帮助信息，描述该参数的用途。

5. parser.add_argument("--model", type=str, default="yolov8n.pt", help="initial weights path")

   • 添加一个命令行参数--model，用于指定初始权重文件的路径。

     • type=str：参数类型为字符串。

     • default="yolov8n.pt"：默认值为"yolov8n.pt"。

     • help="initial weights path"：帮助信息，描述该参数的用途。

6. parser.add_argument("--target-prune-rate", type=float, default=0.2, help="target prune rate")

   • 添加一个命令行参数--target-prune-rate，用于指定目标剪枝率。

     • type=float：参数类型为浮点数。

     • default=0.2：默认值为0.2（即 20%）。

     • help="target prune rate"：帮助信息，描述该参数的用途。

7. parser.add_argument("--iterative-steps", type=int, default=4, help="iterative steps for pruning")

   • 添加一个命令行参数--iterative-steps，用于指定剪枝的迭代步骤数。

     • type=int：参数类型为整数。

     • default=4：默认值为4。

     • help="iterative steps for pruning"：帮助信息，描述该参数的用途。

8. args = parser.parse_args()

   • 解析命令行参数，将解析结果存储在args对象中。

     • parse_args() 是 ArgumentParser 对象的一个方法，用于解析命令行参数。

9. prune(args)

   • 调用prune函数，将解析后的命令行参数传递给该函数，启动剪枝过程。

     • prune 是之前定义的剪枝函数。

代码完整版

import argparse
import math
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "7"
from copy import deepcopy
from datetime import datetime
from pathlib import Path
from typing import List, Union

import numpy as np
import torch
import torch.nn as nn
from matplotlib import pyplot as plt
from ultralytics import YOLO, __version__
from ultralytics.nn.modules import Detect, C2f, Conv, Bottleneck
from ultralytics.nn.tasks import attempt_load_one_weight
from ultralytics.yolo.engine.model import TASK_MAP
from ultralytics.yolo.engine.trainer import BaseTrainer
from ultralytics.yolo.utils import yaml_load, LOGGER, RANK, DEFAULT_CFG_DICT, DEFAULT_CFG_KEYS
from ultralytics.yolo.utils.checks import check_yaml
from ultralytics.yolo.utils.torch_utils import initialize_weights, de_parallel
import torch_pruning as tp

def save_pruning_performance_graph(x, y1, y2, y3):
    # 保存剪枝性能图
    try:
        plt.style.use("ggplot")
    except:
        pass
    x, y1, y2, y3 = np.array(x), np.array(y1), np.array(y2), np.array(y3)
    y2_ratio = y2 / y2[0]
    fig, ax = plt.subplots(figsize=(8, 6))
    ax.set_xlabel('Pruning Ratio')
    ax.set_ylabel('mAP')
    ax.plot(x, y1, label='recovered mAP')
    ax.scatter(x, y1)
    ax.plot(x, y3, color='tab:gray', label='pruned mAP')
    ax.scatter(x, y3, color='tab:gray')
    ax2 = ax.twinx()
    ax2.set_ylabel('MACs')
    ax2.plot(x, y2_ratio, color='tab:orange', label='MACs')
    ax2.scatter(x, y2_ratio, color='tab:orange')
    lines, labels = ax.get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()
    ax2.legend(lines + lines2, labels + labels2, loc='best')
    ax.set_xlim(105, -5)
    ax.set_ylim(0, max(y1) + 0.05)
    ax2.set_ylim(0.05, 1.05)
    max_y1_idx = np.argmax(y1)
    min_y1_idx = np.argmin(y1)
    max_y2_idx = np.argmax(y2)
    min_y2_idx = np.argmin(y2)
    max_y1 = y1[max_y1_idx]
    min_y1 = y1[min_y1_idx]
    max_y2 = y2_ratio[max_y2_idx]
    min_y2 = y2_ratio[min_y2_idx]
    ax.text(x[max_y1_idx], max_y1 - 0.05, f'max mAP = {max_y1:.2f}', fontsize=10)
    ax.text(x[min_y1_idx], min_y1 + 0.02, f'min mAP = {min_y1:.2f}', fontsize=10)
    ax2.text(x[max_y2_idx], max_y2 - 0.05, f'max MACs = {max_y2 * y2[0] / 1e9:.2f}G', fontsize=10)
    ax2.text(x[min_y2_idx], min_y2 + 0.02, f'min MACs = {min_y2 * y2[0] / 1e9:.2f}G', fontsize=10)
    plt.title('Comparison of mAP and MACs with Pruning Ratio')
    plt.savefig('pruning_perf_change.png')

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):
        super().__init__()
        self.c = int(c2 * e)
        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)
        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 = [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()
    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:]
    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:]
    for key in state_dict:
        if not key.startswith('cv1.'):
            state_dict_v2[key] = state_dict[key]
    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):
            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)

def prune(args):
    base_name = 'prune/' + str(datetime.now()) + '/'
    model = YOLO(args.model)
    model.__setattr__("train_v2", train_v2.__get__(model))
    pruning_cfg = yaml_load(check_yaml(args.cfg))
    batch_size = pruning_cfg['batch']
    pruning_cfg['data'] = "./ultralytics/datasets/soccer.yaml"
    pruning_cfg['epochs'] = 4
    model.model.train()
    replace_c2f_with_c2f_v2(model.model)
    initialize_weights(model.model)
    for name, param in model.model.named_parameters():
        param.requires_grad = True
    example_inputs = torch.randn(1, 3, pruning_cfg["imgsz"], pruning_cfg["imgsz"]).to(model.device)
    macs_list, nparams_list, map_list, pruned_map_list = [], [], [], []
    base_macs, base_nparams = tp.utils.count_ops_and_params(model.model, example_inputs)
    pruning_cfg['name'] = base_name + f"baseline_val"
    pruning_cfg['batch'] = 128
    validation_model = deepcopy(model)
    metric = validation_model.val(**pruning_cfg)
    init_map = metric.box.map
    macs_list.append(base_macs)
    nparams_list.append(100)
    map_list.append(init_map)
    pruned_map_list.append(init_map)
    print(f"Before Pruning: MACs={base_macs / 1e9: .5f} G, #Params={base_nparams / 1e6: .5f} M, mAP={init_map: .5f}")
    ch_sparsity = 1 - math.pow((1 - args.target_prune_rate), 1 / args.iterative_steps)
    for i in range(args.iterative_steps):
        model.model.train()
        for name, param in model.model.named_parameters():
            param.requires_grad = True
        ignored_layers = []
        unwrapped_parameters = []
        for m in model.model.modules():
            if isinstance(m, (Detect,)):
                ignored_layers.append(m)
        example_inputs = example_inputs.to(model.device)
        pruner = tp.pruner.GroupNormPruner(model.model, example_inputs,
                                           importance=tp.importance.GroupNormImportance(),
                                           iterative_steps=1,
                                           ch_sparsity=ch_sparsity,
                                           ignored_layers=ignored_layers,
                                           unwrapped_parameters=unwrapped_parameters)
        pruner.step()
        pruning_cfg['name'] = base_name + f"step_{i}_pre_val"
        pruning_cfg['batch'] = 128
        validation_model.model = deepcopy(model.model)
        metric = validation_model.val(**pruning_cfg)
        pruned_map = metric.box.map
        pruned_macs, pruned_nparams = tp.utils.count_ops_and_params(pruner.model, example_inputs.to(model.device))
        current_speed_up = float(macs_list[0]) / pruned_macs
        print(f"After pruning iter {i + 1}: MACs={pruned_macs / 1e9} G, #Params={pruned_nparams / 1e6