datawhale2411组队学习之模型压缩技术1：模型剪枝（上）

项目地址awesome-compression、在线阅读

一、环境配置

整个项目我是在Autodl上跑的。下面是拉取并配置环境的代码：

conda create -n compression python=3.10
conda activate compressiongit clone https://github.com/datawhalechina/awesome-compression.git
cd awesome-compression
pip install -r ./docs/notebook/requirements.txt

测试第三章代码使用plt画图报错没有此字体：

import torch
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D# plt.rcParams['font.sans-serif'] = ['SimHei']  # 解决中文乱码
plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']def plot_tensor(tensor, title):...# 创建一个矩阵weight
weight = torch.rand(8, 8)
plot_tensor(weight, '剪枝前weight')

1. 打印出系统中所有matplotlib已注册的地址及其位置

from matplotlib import font_manager
for font in font_manager.fontManager.ttflist:print(font.name,font.fname)

DejaVu Sans /root/miniconda3/lib/python3.8/site-packages/matplotlib/mpl-data/fonts/ttf/DejaVuSans-Oblique.ttf
DejaVu Sans Mono /root/miniconda3/lib/python3.8/site-packages/matplotlib/mpl-data/fonts/ttf/DejaVuSansMono-Oblique.ttf
...
...

2. 直接将windows10系统字体文件夹C:\Windows\Fonts中的中文字体文件msyh.ttc（微软雅黑）、simfang.ttf（华文仿宋）等上传到系统matplotlib字体文件夹/root/miniconda3/lib/python3.8/site-packages/matplotlib/mpl-data/fonts/ttf下。如果是启用虚拟环境，就放到对应的虚拟环境matplotlib字体文件夹

如果字体文件不在matplotlib字体文件夹，在第三步会报错。

3. 加载中文字体

# 清除字体缓存
font_manager._load_fontmanager(try_read_cache=False)# 下面做主要是根据读取字体文件来获取字体名，即prop.get_name()='Microsoft YaHei'
# 如果你已经直到字体名，直接写plt.rcParams['font.family'] = 'Microsoft YaHei'就行
font_path = '/root/miniconda3/lib/python3.8/site-packages/matplotlib/mpl-data/fonts/ttf/msyh.ttc'  # 替换为你的字体文件路径
prop = font_manager.FontProperties(fname=font_path)
plt.rcParams['font.family'] = prop.get_name()

prop.get_name()用于获取字体名称，此名称和字体文件名不一样。试验了几种字体，微软雅黑效果较好。

二、模型压缩

2.1 模型压缩简介

  下图是近年来模型大小与GPU发展的趋势，从图中可以看出，GPU硬件发展的速度远远跟不上模型大小的增长速度，这也导致了大模型训练和推理的困难。而模型压缩技术可以弥补这个差距，使得大模型可以在有限的硬件资源上运行。

模型压缩技术包括：

1. 模型剪枝：研究发现，神经元可能会出现两种冗余情况：一部分神经元会“坍缩”成功能相似的神经元，共同负责类似的任务，另一些神经元则在优化过程中被忽视。模型剪枝的目标是将这些冗余神经元删除或合并成等效神经元，从而减少模型的参数量和计算量，而不会损害模型的性能，甚至在某些情况下可能会提高性能。

2. 模型量化：神经网络通常使用浮点数进行存储和计算，而减少计算精度（如降低位宽）不会显著影响最终结果。量化技术通过减少模型的数值精度，从而减小存储和计算开销，提高计算速度。

3. 蒸馏学习：将大模型的知识迁移到小模型中，使得小模型能够在不损失性能的情况下，具有类似大模型的表现。

4. 神经网络架构搜索（NAS）：通过自动化搜索优化神经网络架构，找到适合的、较小的网络结构。

2.2 模型压缩评价指标

指标	描述	备注
准确率 (Accuracy)	对比模型压缩前后在特定任务上的准确度，如分类准确率、检测精度等。	通常计算在零样本（Zero-shot）数据集上的准确率
参数量 (Params)	模型中可训练参数的总数，包括所有权重和偏置。	模型大小通常是通过模型的总参数数量来衡量的，参数越多计算资源和内存消耗也越多
模型大小 (Model Size)	以存储大小（如MB）衡量压缩效果，计算公式为：大小 = 参数量 * 带宽。	一个模型的参数量为70B，假设使用32位浮点数存储，那么其模型大小为：$70B\ast 4Bytes(32bits)=280GB(280\ast 10^{9}Bytes)$
乘累加操作 (MACs)	浮点运算的基本单位，包括一个乘法操作和一个加法操作，可描述CNN卷积操作的计算量。	卷积层前向传播时，每个卷积核都会与输入特征图进行点积，然后累加所有点积结果，得到输出特征图。此过程就是一次MAC操作。
浮点运算 (FLOPs，Number of Floating Point Operations)	模型执行一次前向推理需要的浮点运算次数。	代表LLM执行一个实例所需的浮点运算数量，与模型需要计算资源成正比
操作数 (OPs)	神经网络中的激活或权重计算也不总是浮点运算，OPS(Operation Per Second)代表每秒执行的操作数
压缩比 (Compression Ratio)	原始模型大小与压缩后模型大小的比值，越高压缩效果越好，但可能伴随性能损失。
推理时间 (Inference Time)	模型处理输入并生成输出的时间。	压缩后通常推理时间会减少，提升响应速度。
吞吐量 (Throughput)	模型单位时间内处理的数据量，衡量压缩后模型的效率。	通常与准确率和延迟一起考量，平衡模型精度和推理速度之间的关系

三、 模型剪枝

3.1 模型剪枝简介

• 概念：模型剪枝通过移除模型中不重要的权重（即神经网络中的突触，Synapses）和分支（网络中的神经元，Neurons），将网络结构稀疏化，减少模型的参数量和计算复杂度，从而降低内存消耗和加速模型推理过程。这种技术对于在资源受限的设备上部署模型尤为重要。剪枝过程如下图所示：
  

• 目标：在减少模型大小的同时，尽量保持模型性能，找到模型大小和性能之间的最佳平衡点。

• 数学表示：对于一个简单的线性神经网络，其公式可表示为：
  $$Y=WX$$
    其中，$W$为权重，$X$为输入，即神经元。模型剪枝可以看作是将权重矩阵中的一些元素设置为零，这些剪枝后具有大量零元素的矩阵被称为稀疏矩阵，反之绝大部分元素非零的矩阵被称为稠密矩阵。

剪枝的理论基础：

• 彩票假说：大型神经网络中存在一个子网络，如果独立训练，可以达到与完整网络相似的性能。
• 网络稀疏性：许多深度神经网络参数呈现出稀疏性，即大部分参数值接近于零。这种稀疏性启发了剪枝技术，即通过移除这些非显著的参数来简化模型。
• 正则化：L1正则化鼓励网络学习稀疏的参数分布，这些参数权重接近于零，可以安全移除。

3.2 何为剪枝（What is Pruning?）

剪枝可以按照不同标准进行划分，下面主要从剪枝类型、剪枝范围、剪枝粒度三方面来阐述。

• 剪枝类型：

  • 非结构化剪枝：不关心权重在网络中的位置，仅根据权重的大小决定是否剪枝（剪掉较小的权重），导致模型稀疏。这种方法会破坏原有模型的结构，不利于现有硬件加速，需要设计特定的硬件。
  • 结构化剪枝：考虑了网络的结构，通常是移除整个神经元、卷积核或层，易于硬件加速。结构化剪枝的实现通常更加复杂，需要更多的计算资源和对模型的深入理解
  • 半结构化剪枝：两者之间的折中，在剪枝时部分考虑结构，但不完全按整体结构剪去某个单元，试图在保持计算加速的同时减少性能损。比如移除整个神经元或过滤器的一部分，而不是全部。

    NVIDIA A100 GPU在稀疏加速方面的一个显著特性是其支持的2:4 （50%）稀疏模式，这是一种特定形式的半结构化剪枝。在这种模式下，模型中的每个权重块（通常是 4 个连续的权重值）中有 2 个权重为非零值，剩下 2 个权重置为0。

• 剪枝范围：

  • 局部剪枝：局部剪枝的操作是独立进行的，不依赖于模型的整体结构，直接移除那些对模型输出影响较小的权重。可以是权重剪枝，也可以是神经元剪枝，甚至是通道剪枝（CNN中）。
  • 全局剪枝：考虑模型的整体结构和性能，可能会移除整个神经元、卷积核、层或更复杂的结构，通常需要对模型有深入的理解，且可能涉及重设计模型架构。

• 剪枝粒度：按照剪枝粒度进行划分，剪枝可分为：

  • 细粒度剪枝（Fine-grained Pruning）：移除权重矩阵中的任意值，可以实现高压缩比，但对硬件不友好，因此速度增益有限。
  • 基于模式的剪枝（Pattern-based Pruning）：在第四章剪枝实践中进行详细介绍
  • 向量级剪枝（Vector-level Pruning）：以行或列为单位对权重进行裁剪。
  • 内核级剪枝（Kernel-level Pruning）：卷积核（滤波器）为单位对权重进行裁剪。
  • 通道级剪枝（Channel-level Pruning）：以通道为单位对权重进行裁剪。
    

  通道级剪枝是改变了网络中的滤波器组和特征通道数目，属于结构化剪枝。其它四者使网络的拓扑结构本身发生了变化，需要专门的算法设计来支持这种稀疏的运算，属于非结构化剪枝。

3.3 剪枝标准（How to prune?）

  剪枝算法通常基于权重的重要性来决定是否剪枝。权重的重要性可以通过多种方式评估，例如权重的大小、权重对损失函数的梯度、或者权重对输入的激活情况等。

1. 基于权重大小：直接根据每个元素的权重绝对值大小来计算权重的重要性，移除绝对值较小的权重，整个剪枝过程如下：
   
   此外，也可以根据权重的L1和L2正则化来指导指导模型剪枝。具体来说，以行为单位，计算每行的重要性，移除权重中那些重要性较小的行
   
   

2. 基于梯度大小：

   • 以权值大小为依据的进行剪枝很容易剪掉重要的权值。以人脸识别为例，在人脸的诸多特征中，眼睛的细微变化如颜色、大小、形状，对于人脸识别的结果有很大影响，不应该被剪掉。
   • 在模型训练过程中，梯度是损失函数对权值的偏导数，反映了损失对权值的敏感程度，梯度越大的权重越重要，所以应该去除较小梯度的权重。

3. 基于尺度(Scaling-based)：利用BN层中的缩放因子来实现稀疏性，识别并剪枝那些对模型输出影响不大的整个通道。
   

在标准的CNN训练中，批归一化（BN）层计算公式如下：
$$\hat{x}=\frac{x-\mu }{\sqrt{{\sigma }^2+ϵ}}$$
其中，$\mu$ 是批次均值，${\sigma }^2$ 是批次方差，$ϵ$ 是为了数值稳定而加的一个小值。
在此基础上，BN层会应用一个可训练的缩放因子 $\gamma$(调节标准化后数据的尺度，影响输出的幅度) 和一个可训练的偏移量 $\beta$(平移标准化后数据的分布)：
$$y=\gamma \hat{x}+\beta$$
Network Slimming 方法中，在损失函数中添加一个L1正则化项来鼓励 $\gamma$趋向于零，从而可以识别不重要的通道并进行剪枝。

4. 基于二阶梯度（Second-Order-based）：最具代表性的是最优脑损伤（Optimal Brain Damage，OBD），它使用 Hessian 矩阵（损失函数对网络权重的二阶导数矩阵）来判断每个权重的重要性。

   • 使用条件：假设神经网络的损失函数可以在一个局部区域内近似为一个 二次函数；神经网络的训练已收敛； 删除每个参数所导致的误差是独立的。
   • 误差计算：当某个权重 $w_i$ 被移除时，模型的损失函数变化可用公式近似：$\delta L_i\approx \frac{1}{2}{h}_{ii}{w}_i^2$，其中，$h_{ii}$ 是 Hessian 矩阵的对角元素，表示损失函数关于权重 $w_i$ 的二阶导数，即$h_{ii}=\frac{{\partial }^{2}L}{\partial w_i^2}$
   • 剪枝原则：误差公式表明：权重的重要性由其对应的 Hessian 对角线元素和其平方决定。对于较小的 $|\delta L_i|$，权重的重要性较低，可以被剪去。

OBD剪枝的优点有：

• 评估更准确：相比基于一阶梯度（如 L1 或 L2 正则化）方法提供的梯度方向和大小，二阶导数考虑了梯度的变化速率，而不仅依赖于当前梯度，这使得它能够更精确地评估权重的重要性。
• 适合收敛后的模型：在模型已经收敛的情况下，一阶梯度可能接近零，但权重的二阶效应（如曲率）仍能反映权重的重要性。

OBD剪枝的缺点是：

• 计算复杂度高：纵使OBD 只需要考虑 Hessian 矩阵的对角线元素 $h_{ii}$，而无需考虑整个 Hessian 矩阵，但仍然比计算一阶梯度复杂。
• 近似性限制：OBD 假设损失函数在当前权重点附近是二次的，这在一些复杂的非线性模型中可能并不准确。
• 局限于收敛模型：OBD 的假设前提之一是模型已经收敛，对于训练初期的网络，二阶信息的作用可能较小。

  总结来说，OBD 使用 Hessian 矩阵来判断权重的重要性是为了利用其提供的曲率信息，进行更精确的剪枝评估。虽然计算复杂度较高，但能避免简单一阶梯度方法的局限性，适用于在模型已经收敛的情况下进行精细剪枝。

3.4 剪枝频率（How often?）

• 迭代剪枝（Iterative Pruning）：逐步移除权重，可以更细致地评估每一次剪枝对模型性能的影响，并允许模型有机会调整其余权重来补偿被剪除的权重。
  • 训练模型：首先训练一个完整的、未剪枝的模型，使其在训练数据上达到一个良好的性能水平。
  • 剪枝：使用一个预定的剪枝策略来轻微剪枝网络，移除一小部分权重。
  • 微调：使用原始训练数据集重新训练模型，以恢复由于剪枝引起的性能损失。
  • 评估：在验证集上评估剪枝后模型的性能，确保模型仍然能够维持良好的性能。
  • 重复：重复步骤2到步骤4，每次迭代剪掉更多的权重，并进行微调，直到达到一个预定的性能标准或剪枝比例。
• 单次剪枝（One-Shot Pruning）：模型在被训练到收敛后，对其进行一次性的剪枝操作，优点是高效直接，缺点是会极大地受到噪声的影响。

  迭代式剪枝在每次迭代之后只会删除掉少量的权重，然后周而复始地进行其他轮的评估和删除，能够在一定程度上减少噪声对于整个剪枝过程的影响。但对于大模型来说，由于微调的成本太高，所以更倾向于使用单次剪枝方法。

3.5 剪枝时机（When to prune?）

• 训练后剪枝：先训练一个模型，然后对模型进行剪枝，最后对剪枝后模型进行微调。
• 训练时剪枝：直接在模型训练过程中进行剪枝，最后对剪枝后模型进行微调。
• 训练前剪枝：在模型训练前进行剪枝，然后从头训练剪枝后的模型。

3.6 剪枝比例

• 均匀分层剪枝（Uniform Layer-Wise Pruning）：在神经网络的每一层中都应用相同的剪枝率。这种方法实现简单，剪枝率容易控制，但它忽略了每一层对模型整体性能的重要性差异。
• 非均匀分层剪枝（Non-Uniform Layer-Wise Pruning）：根据每一层的不同特点来分配不同的剪枝率。例如，可以根据梯度信息、权重的大小、或者其他指标（如信息熵、Hessian矩阵等）来确定每一层的剪枝率，越重要的层，保留的参数越多。非均匀剪枝往往比均匀剪枝的性能更好。
  

四、代码实践

原始代码见awesome-compression/docs/notebook/ch03

下面先讲点前置知识，方便后面剪枝粒度的理解。

• Kernel：卷积核，卷积层中用于特征提取的小矩阵，其size为（height,width)
• Channel：通道，通常指数据的深度维度。例如，彩色图像有RGB三个颜色通道。在CNN中，输入层的Channel数对应于图像的颜色通道数，而隐藏层的Channel数则对应于该层Filter的数量，即每个Filter产生的特征图数量。
• Filter：滤波器。在对多通道输入进行卷积时，每个Channel都有其对应的Kernel，它们的集合构成了一个Filter。例如，对于一个RGB彩色图像，一个Filter将包含三个卷积核，每个用于一个颜色通道。最终输出是一个二维的特征图（Feature Map）。Filter能够检测特定类型的特征，不同的Filter可以捕捉到不同的特征。
• Feature Map：特征图，指的是从输入数据（如图像）中通过特定的卷积滤波器（Filter）提取出的特征表示。
• Layer：CNN由多个层组成，每个层可以是卷积层、池化层、全连接层等。每个卷积层由多个Filter组成，每个Filter通过卷积操作生成一个特征图，所有特征图堆叠在一起形成该层的输出。
• Kernel/Filter关注的是局部特征的提取，Channel关注的是特征的多样性和表示，而Layer则是网络结构的组成部分

下图是对一个3通道的图片做卷积操作，即对三通道图片使用一个Filter：

如果是对三通道图片使用两个Filter，结果就是两个特征图：

4.1 剪枝粒度实践

import torch
import matplotlib.pyplot as plt
from matplotlib import font_manager
from mpl_toolkits.mplot3d import Axes3D#font_manager._load_fontmanager(try_read_cache=False)
plt.rcParams['font.family'] = 'Microsoft YaHei'

# 创建一个可视化2维矩阵函数，将值为0的元素与其他区分开
def plot_tensor(tensor, title):# 创建一个新的图像和和一组子图轴ax（subplots函数返回一个图像对象和一个轴对象）fig, ax = plt.subplots()# tensor.cpu().numpy()将张量转为numpy数组，ax.imshow(tensor.cpu().numpy() == 0, vmin=0, vmax=1, cmap='tab20c')ax.set_title(title)ax.set_yticklabels([])ax.set_xticklabels([])# 遍历矩阵中的每个元素并添加文本标签for i in range(tensor.shape[1]):for j in range(tensor.shape[0]):text = ax.text(j, i, f'{tensor[i, j].item():.2f}', ha="center", va="center", color="k")# 显示图像plt.show()

• cmap='tab20c'是一个包含20种颜色颜色映射方案，vmin和vmax参数设置颜色映射的值范围（超过此范围的数值被裁剪到最大值或最小值）。在tab20c颜色映射中，True值（1）通常会被映射到颜色映射中的一个颜色，而False值（0）会被映射到另一种颜色。
• text = ax.text(j, i, f'{tensor[i, j].item():.2f}', ha="center", va="center", color="k")：在每个元素的位置添加文本标签，显示该元素的值（保留两位小数）。ha="center",va="center"参数设置文本的水平和垂直对齐方式为居中。color="k"设置文本颜色为黑色。

4.1.1 细粒度剪枝

# 1. 创建一个随机矩阵weight
weight = torch.rand(8, 8)
plot_tensor(weight, '剪枝前weight')# 2. 随机定义一个剪枝规则，比如将Tensor里的值小于0.5的都置为0
def _fine_grained_prune(tensor: torch.Tensor, threshold  : float) -> torch.Tensor:""":param tensor: 输入张量，包含需要剪枝的权重。:param threshold: 阈值，用于判断权重的大小。:return: 剪枝后的张量。"""for i in range(tensor.shape[1]):for j in range(tensor.shape[0]):if tensor[i, j] < threshold:tensor[i][j] = 0return tensorpruned_weight = _fine_grained_prune(weight, 0.5)
plot_tensor(weight, '细粒度剪枝后weight')

  用for循环遍历去实现，虽然结果是对的，但如果参数太大的话，肯定会影响到速度。剪枝中常用的一种方法是使用mask掩码矩阵来实现。

# 3. 使用mask矩阵实现细粒度剪枝
def fine_grained_prune(tensor: torch.Tensor, threshold  : float) -> torch.Tensor:"""创建一个掩码张量，指示哪些权重不应被剪枝（应保持非零）。:param tensor: 输入张量，包含需要剪枝的权重。:param threshold: 阈值，用于判断权重的大小。:return: 剪枝后的张量。"""mask = torch.gt(tensor, threshold)        # 所有大于threshold的位置返回Truetensor.mul_(mask)						  # mul_表示原地更新的乘法操作return tensor
pruned_weight = fine_grained_prune(weight, 0.5)
plot_tensor(pruned_weight, '细粒度剪枝后weight')

4.1.2 基于模式的剪枝

  N:M 稀疏度表示 DNN 的稀疏度，即每M个连续权重中固定有N个非零值，其余元素均置为0。这种结构可以利用NVIDIA的稀疏张量核心加速矩阵乘法，比如NVIDIA Ampere A100 GPU通过稀疏张量核心支持2:4稀疏度，实现高达2倍的吞吐量提升，而不影响计算的准确性。

  稀疏矩阵W首先会被压缩，压缩后的矩阵存储着非零的数据值，而metadata则存储着对应非零元素在原矩阵W中的索引信息（非零元素的行号和列号压缩成两个独立的一维数组）如下图所示：

  如下图所示，以NVIDIA 2：4为例：

• 创建一个patterns，从4个中取出2个为非零值，可以算出一共有6种不同的模式（代表每行有6种剪枝方式）
• 将weight matrix变换成nx4的格式方便与pattern进行矩阵运算，运算后的结果为nx6的矩阵
• 在n的维度上进行argmax取得最大的索引，表示这一行剪枝的最佳方式，将索引对应的pattern值填充到mask中。
• 将mask reshape到weight matrix相同的形式，二者进行矩阵乘法，得到剪枝后的结果

# 创建一个矩阵weight
weight = torch.rand(8, 8)
plot_tensor(weight, '剪枝前weight')

from itertools import permutationsdef reshape_1d(tensor, m):"""将输入的二维张量 tensor 重新塑形为列数为 m 的矩阵。如果原始张量的列数不能被 m 整除，它会填充零以使列数成为 m 的倍数。"""if tensor.shape[1] % m > 0:mat = torch.FloatTensor(tensor.shape[0], tensor.shape[1] + (m - tensor.shape[1] % m)).fill_(0)mat[:, : tensor.shape[1]] = tensorreturn mat.view(-1, m)else:return tensor.view(-1, m)def compute_valid_1d_patterns(m, n):# 创建一个长度为 m 的全零张量，然后将前 n 个元素设置为 1patterns = torch.zeros(m)patterns[:n] = 1# 使用 permutations 生成所有可能的排列，并转换为 torch.Tensor。valid_patterns = torch.Tensor(list(set(permutations(patterns.tolist()))))# 最终返回6行4列的矩阵，每行是一种4：2的模式return valid_patternsdef compute_mask(tensor, m, n):# 计算所有可能的模式patterns = compute_valid_1d_patterns(m,n)# 初始化一个与tensor形状相同的全1张量，然后reshape到m列(8*8 ——> 16*4)mask = torch.IntTensor(tensor.shape).fill_(1).view(-1,m)mat = reshape_1d(tensor, m)# 计算mat的绝对值与所有可能模式的转置（patterns.t()）的矩阵乘积，然后得到最佳索引pmax = torch.argmax(torch.matmul(mat.abs(), patterns.t()), dim=1)# 从所有可能的模式中选择最佳模式，并更新 mask 张量。这样mask 中的每个元素将反映输入张量中相应元素的最佳模式。mask[:] = patterns[pmax[:]]# 将 mask 张量的形状调整回与输入张量 tensor 相同的形状，以便可以直接应用于原始张量。mask = mask.view(tensor.shape)return maskdef pattern_pruning(tensor, m, n):mask = compute_mask(tensor, m, n)tensor.mul_(mask)return tensorpruned_weight = pattern_pruning(weight, 4, 2)
plot_tensor(pruned_weight, '剪枝后weight')

4.1.3 向量级别剪枝

# 创建一个矩阵weight
weight = torch.rand(8, 8)
plot_tensor(weight, '剪枝前weight')# 剪枝某个点所在的行与列
def vector_pruning(weight, point):row, col = pointprune_weight = weight.clone()prune_weight[row, :] = 0prune_weight[:, col] = 0return prune_weight
point = (1, 1)
prune_weight = vector_pruning(weight, point)
plot_tensor(prune_weight, '向量级剪枝后weight')

4.1.4 卷积核级别剪枝

# 定义可视化4维张量的函数
def visualize_tensor(tensor, title, batch_spacing=3):fig = plt.figure()  # 创建一个新的matplotlib图形ax = fig.add_subplot(111, projection='3d')  # 向图形中添加一个3D子图# 遍历张量的批次维度for batch in range(tensor.shape[0]):# 遍历张量的通道维度for channel in range(tensor.shape[1]):# 遍历张量的高度维度for i in range(tensor.shape[2]):# 遍历张量的宽度维度for j in range(tensor.shape[3]):# 计算条形的x位置，考虑到不同批次间的间隔x = j + (batch * (tensor.shape[3] + batch_spacing))y = i  # 条形的y位置，即张量的高度维度z = channel  # 条形的z位置，即张量的通道维度# 如果张量在当前位置的值为0，则设置条形颜色为红色，否则为绿色color = 'red' if tensor[batch, channel, i, j] == 0 else 'green'# 绘制单个3D条形ax.bar3d(x, y, z, 1, 1, 1, shade=True, color=color, edgecolor='black', alpha=0.9)ax.set_title(title)  # 设置3D图形的标题ax.set_xlabel('Width')  # 设置x轴标签，对应张量的宽度维度ax.set_ylabel('Height')  # 设置y轴标签，对应张量的高度维度ax.set_zlabel('Channel')  # 设置z轴标签，对于张量的通道维度ax.set_zlim(ax.get_zlim()[::-1])  # 反转z轴方向ax.zaxis.labelpad = 15  # 调整z轴标签的填充plt.show()  # 显示图形

def prune_conv_layer(conv_layer,title="", percentile=0.2, vis=True, dim=None):"""conv_layer: 输入的卷积层（或张量）。title: 用于可视化的标题，默认为空字符串。percentile: 用于确定剪枝阈值的百分位数，默认为 0.2。vis: 是否进行可视化，默认为 True。dim: 计算 L2 范数的维度，默认为 None，如果指定，将在这个维度上计算 L2 范数"""prune_layer = conv_layer.clone()l2_norm = Nonemask = None# 计算每个kernel的L2范数,keepdim=True 保持输出的维度与输入一致。l2_norm = torch.norm(prune_layer, p=2, dim=dim, keepdim=True)# 计算L2范数的20%位数threshold = torch.quantile(l2_norm, percentile)mask = l2_norm > thresholdprune_layer = prune_layer * mask.float()visualize_tensor(prune_layer,title=title)  

   torch.quantile(input, q) 是 PyTorch 提供的一个函数，用于计算输入张量input在给定百分位数q处的值。比如torch.quantile(tensor, 0.5)表示中位数，torch.quantile(tensor, 0.75)表示75%位数。

下面使用PyTorch创建一个四维张量(N, C, H, W)，每个维度分别代表(batch_size,channel,height,width)，然后进行kernel级别剪枝：

tensor = torch.rand((3, 10, 4, 5))
# 调用函数进行剪枝，dim=(-2, -1)表示高和宽两个维度
pruned_tensor = prune_conv_layer(conv_layer=tensor, title='Kernel级别剪枝',dim=(2, 3), vis=True)

4.1.5 通道级别剪枝

pruned_tensor = prune_conv_layer(conv_layer=tensor, title='Channel级别剪枝',dim=(0, 2, 3), vis=True)

4.1.6 滤波器级别剪枝

pruned_tensor = prune_conv_layer(conv_layer=tensor, title='Filter级别剪枝',dim=(1, 2, 3), vis=True)

4.1.7 汇总

# 返回一个mask
def get_threshold_and_mask(norms, percentile):threshold = torch.quantile(norms, percentile)return norms > thresholddef prune_conv_layer(conv_layer, prune_method, title= "", percentile=0.2, vis=True):prune_layer = conv_layer.clone()mask = Noneif prune_method == "fine_grained":prune_layer[torch.abs(prune_layer) < percentile] = 0elif prune_method == "vector_level":mask = get_threshold_and_mask(torch.norm(prune_layer, p=2, dim=-1), percentile).unsqueeze(-1)elif prune_method == "kernel_level":mask = get_threshold_and_mask(torch.norm(prune_layer, p=2, dim=(-2, -1), keepdim=True), percentile)elif prune_method == "filter_level":mask = get_threshold_and_mask(torch.norm(prune_layer, p=2, dim=(1, 2, 3), keepdim=True), percentile)elif prune_method == "channel_level":mask = get_threshold_and_mask(torch.norm(prune_layer, p=2, dim=(0, 2, 3), keepdim=True), percentile)if mask is not None:prune_layer = prune_layer * mask.float()if vis:visualize_tensor(prune_layer, title=title)  # 实现可视化的函数return prune_layer# 使用PyTorch创建一个张量
tensor = torch.rand((3, 10, 4, 5)) # 调用函数进行剪枝
pruned_tensor = prune_conv_layer(tensor, 'fine_grained', '细粒度剪枝',  vis=True)
pruned_tensor = prune_conv_layer(tensor, 'vector_level', 'Vector级别剪枝', vis=True)
pruned_tensor = prune_conv_layer(tensor, 'kernel_level', 'Kernel级别剪枝', vis=True)
pruned_tensor = prune_conv_layer(tensor, 'filter_level', 'Filter级别剪枝', vis=True)
pruned_tensor = prune_conv_layer(tensor, 'channel_level', 'Channel级别剪枝', vis=True)

4.2 剪枝标准实践

4.2.1 定义初始网络，画出权重分布图和密度直方图

import copy
import math
import random
import timeimport torch
import torch.nn as nn
import numpy as np
from matplotlib import pyplot as plt
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision import datasets
import torch.nn.functional as F# 设置 matplotlib 使用支持负号的字体
plt.rcParams['font.family'] = 'DejaVu Sans'# 1. 定义一个LeNet网络
class LeNet(nn.Module):def __init__(self, num_classes=10):super(LeNet, self).__init__()self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5)self.conv2 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5)self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)self.fc1 = nn.Linear(in_features=16 * 4 * 4, out_features=120)self.fc2 = nn.Linear(in_features=120, out_features=84)self.fc3 = nn.Linear(in_features=84, out_features=num_classes)def forward(self, x):x = self.maxpool(F.relu(self.conv1(x)))x = self.maxpool(F.relu(self.conv2(x)))x = x.view(x.size()[0], -1)x = F.relu(self.fc1(x))x = F.relu(self.fc2(x))x = self.fc3(x)return x
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = LeNet().to(device=device)
model

LeNet((conv1): Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1))(conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))(maxpool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)(fc1): Linear(in_features=256, out_features=120, bias=True)(fc2): Linear(in_features=120, out_features=84, bias=True)(fc3): Linear(in_features=84, out_features=10, bias=True)
)

# 2. 加载模型
checkpoint = torch.load('../ch02/model.pt')
# 加载状态字典到模型
model.load_state_dict(checkpoint)
origin_model = copy.deepcopy(model)

# 3. 绘制权重分布图
def plot_weight_distribution(model, bins=256, count_nonzero_only=False):fig, axes = plt.subplots(2,3, figsize=(10, 6))		# 创建了一个2行3列的子图布局，整个图形的大小为10x6英寸    fig.delaxes(axes[1][2])								# 实际只绘制5个子图，所以我们删除第6个子图    axes = axes.ravel()									# 将多维数组axes展平成一维数组，方便后续索引plot_index = 0										# 初始化一个索引变量plot_index，用于跟踪当前正在绘制的子图# 遍历模型的所有参数（参数名称和参数本身的迭代器）for name, param in model.named_parameters():if param.dim() > 1:ax = axes[plot_index]# 如果count_nonzero_only为True，则只考虑非零权重if count_nonzero_only:# 将参数从计算图中分离出来，展平成一维数组，并移动到CPU上param_cpu = param.detach().view(-1).cpu()# 筛选出非零的权重param_cpu = param_cpu[param_cpu != 0].view(-1)# 在当前子图上绘制直方图，bins指定柱状区间的数量# density=True表示显示概率密度，color和alpha分别设置柱状图的颜色和透明度。ax.hist(param_cpu, bins=bins, density=True, color = 'green', alpha = 0.5)else:# 如果count_nonzero_only为False，则绘制所有权重的直方图ax.hist(param.detach().view(-1).cpu(), bins=bins, density=True, color = 'green', alpha = 0.5)ax.set_xlabel(name)ax.set_ylabel('density')plot_index += 1# 设置整个图形的标题，并自动调整子图参数，使之填充整个图形区域。fig.suptitle('Histogram of Weights')fig.tight_layout()fig.subplots_adjust(top=0.925)plt.show()plot_weight_distribution(model)

  最终显示的是模型各层权重的分布情况，x 轴是权重值，y 轴是概率密度。这些图也有助于理解模型的权重在训练过程中的情况，例如是否有梯度消失（权重值集中在零附近）或爆炸（权重分布会出现极大值）的问题。

# 4. 计算每一层网络的稠密程度（非0参数的占比）
def plot_num_parameters_distribution(model):# 创建一个空字典，用于存储每一层的名称和对应的参数密度。num_parameters = dict()# 初始化两个变量，分别用于存储每一层中非零参数的数量和总参数数量num_nonzeros, num_elements = 0, 0for name,param in model.named_parameters():if param.dim() > 1:# 计算当前参数中非零元素的数量和总数量num_nonzeros = param.count_nonzero()num_elements = param.numel()dense = float(num_nonzeros) / num_elementsnum_parameters[name] = dense#创建一个大小为8x6英寸的图形，在图形上添加y轴网格线fig = plt.figure(figsize=(8, 6))plt.grid(axis='y')# 使用柱状图显示每一层的参数密度，x轴为参数名称，y轴为参数密度。bars = plt.bar(list(num_parameters.keys()), list(num_parameters.values()))# 在柱状图上添加数据标签for bar in bars:# 获取柱状图的高度值。yval = bar.get_height()# 在每个柱状图上方添加文本标签，显示柱状图的高度值。plt.text(bar.get_x() + bar.get_width()/2.0, yval, yval, va='bottom')  plt.title('#Parameter Distribution')plt.ylabel('Number of Parameters')# 将x轴的标签旋转60度，以便更好地显示参数名称plt.xticks(rotation=60)# 自动调整子图参数，使之填充整个图形区域。plt.tight_layout()plt.show()# 绘制模型中每一层参数密度的分布柱状图
plot_num_parameters_distribution(model)

4.2.2 基于L1权重大小的剪枝

@torch.no_grad()
def prune_l1(weight, percentile=0.5):num_elements = weight.numel()									# 权重总数   num_zeros = round(num_elements * percentile)					# 将50%的权重置为0    importance = weight.abs()										# 计算权重的绝对值，作为权重的重要性指标			    threshold = importance.view(-1).kthvalue(num_zeros).values		#  计算裁剪阈值    																mask = torch.gt(importance, threshold)							# 创建mask，大于阈值的权重为True	       weight.mul_(mask)												# 计算mask后的weightreturn weight

  importance.view(-1).kthvalue(num_zeros).values：先将importance展平为一维向量，然后求其第k小的元素（k=num_zeros）的值（values）。

weight_pruned = prune_l1(model.conv2.weight, percentile=0.5)		# 裁剪conv2层
model.conv2.weight.data = weight_pruned								# 替换原有model层
plot_weight_distribution(model)										# 画出weight直方图

# 画出weight稠密度直方图
plot_num_parameters_distribution(model)

4.2.3 基于L2权重大小的剪枝

@torch.no_grad()
def prune_l2(weight, percentile=0.5):num_elements = weight.numel()    num_zeros = round(num_elements * percentile)					# 计算值为0的数量    importance = weight.pow(2)										# 计算weight的重要性（使用L2范数，即各元素的平方）   threshold = importance.view(-1).kthvalue(num_zeros).values		# 计算裁剪阈值    mask = torch.gt(importance, threshold)							# 计算mask        weight.mul_(mask)												# 计算mask后的weightreturn weight

weight_pruned = prune_l2(model.fc1.weight, percentile=0.4)			# 裁剪fc1层
model.fc1.weight.data = weight_pruned								# 替换原有model层
plot_weight_distribution(model										# 画出weight直方图

# 画出weight稠密度直方图
plot_num_parameters_distribution(model)
# 保存裁剪后的weight
torch.save(model.state_dict(), './model_pruned.pt')

4.2.4 基于梯度大小的剪枝

del model
model = LeNet().to(device=device)gradients = torch.load('../ch02/model_gradients.pt')		# 加载梯度信息
checkpoint = torch.load('../ch02/model.pt')					# 加载参数信息
model.load_state_dict(checkpoint)							# 加载状态字典到模型

@torch.no_grad()
def gradient_magnitude_pruning(weight, gradient, percentile=0.5):num_elements = weight.numel()    num_zeros = round(num_elements * percentile)					# 计算值为0的数量    importance = gradient.abs()										# 计算weight的重要性（使用L1范数）    threshold = importance.view(-1).kthvalue(num_zeros).values		# 计算裁剪阈值    mask = torch.gt(importance, threshold)							# 计算mask    weight.mul_(mask)												# 计算mask后的weightreturn weight

# 对fc2应用梯度裁剪
gradient_magnitude_pruning(model.fc2.weight, gradients['fc2.weight'], percentile=0.5)
# 列出weight直方图
plot_weight_distribution(model)

plot_num_parameters_distribution(model)

4.3 剪枝时机实践（训练后剪枝）

模型训练后进行剪枝的步骤如下：

1. 初始训练：使用反向传播算法训练神经网络，学习权重和网络结构。
2. 识别重要连接：训练后，识别出对输出影响显著的连接，通常是权重较大的连接。
3. 设置阈值：选择一个阈值，低于该阈值的连接被视为不重要。
4. 剪枝：移除所有权重低于阈值的连接，形成稀疏层。
5. 重新训练：剪枝后重新训练网络，调整剩余连接的权重以保持准确性。
6. 迭代剪枝：重复剪枝和重新训练的过程，直到达到在不显著损失准确性的情况下尽可能减少连接的平衡点。

  下面还是以上一章节定义的LeNet网络为例，根据权重绝对值大小进行剪枝，并评估其剪枝前后在mnist数据集上的性能，以及模型指标。

4.3.1 加载LeNet网络，评估其权重分布和模型指标

1. 定义一些神经网络模型的的属性函数，包括模型的MACs（乘累加操作数），稀疏度，参数总数和模型大小。

from torchprofile import profile_macs# 1. 定义一个函数，用于获取模型的MACs（乘累加操作数）。
def get_model_macs(model, inputs) -> int:return profile_macs(model, inputs)# 2. 定义一个函数，用于计算给定张量的稀疏度，即张量中零元素的占比。
def get_sparsity(tensor: torch.Tensor) -> float:# 计算张量中非零元素的数量num_nonzeros = tensor.count_nonzero()# 计算张量中总元素的数量num_elements = tensor.numel()# 计算稀疏度并返回return 1 - float(num_nonzeros) / num_elements# 3. 定义一个函数，用于计算给定模型的稀疏度。
def get_model_sparsity(model: nn.Module) -> float:num_nonzeros, num_elements = 0, 0for param in model.parameters():num_nonzeros += param.count_nonzero()num_elements += param.numel()return 1 - float(num_nonzeros) / num_elements# 4. 定义一个函数，用于计算模型的总参数数量。
def get_num_parameters(model: nn.Module, count_nonzero_only=False) -> int:# 参数count_nonzero_only决定是否只计算非零权重的参数。num_counted_elements = 0for param in model.parameters():if count_nonzero_only:num_counted_elements += param.count_nonzero()else:num_counted_elements += param.numel()return num_counted_elements# 5. 定义一个函数，用于计算模型的大小，单位为比特。
def get_model_size(model: nn.Module, data_width=32, count_nonzero_only=False) -> int:# 参数data_width表示每个元素的位数，count_nonzero_only决定是否只计算非零权重。return get_num_parameters(model, count_nonzero_only) * data_width# 6. 定义一些常用的数据单位常量
Byte = 8
KiB = 1024 * Byte
MiB = 1024 * KiB
GiB = 1024 * MiB

2. 数据预处理（略）
   • 定义transform 方式
   • 使用datasets.MNIST加载mnist训练集和测试集数据
   • 将数据装入train_loader 和test_loader
3. 定义训练函数train和评估函数evaluate（略）
4. 定义一个LeNet网络，加载模型。初始模型的大小和准确率为：

# 5. 加载训练后的模型
model = LeNet()
checkpoint = torch.load('../ch02/model.pt')
# 加载状态字典到模型
model.load_state_dict(checkpoint)# 6. 备份并评估model
origin_model = copy.deepcopy(model)
origin_model_accuracy = evaluate(origin_model, test_loader)
origin_model_size = get_model_size(origin_model)
print(f"dense model has accuracy={origin_model_accuracy:.2f}%")
print(f"dense model has size={origin_model_size/MiB:.2f}%MiB")

eval:   0%|          | 0/157 [00:00<?, ?it/s]
dense model has accuracy=97.99%
dense model has size=0.17%MiB

# 7. 绘制weight分布图和稀疏度直方图
plot_weight_distribution(model)
plot_num_parameters_distribution(model)

4.3.2 模型剪枝

1. 定义细粒度剪枝函数fine_grained_prune，对单个张量进行基于绝对值的剪枝。
   • 定义剪枝比例sparsity
   • 根据权重的绝对值importance = tensor.abs()来确定每个权重的重要性。
   • 根据剪枝比例计算阈值threshold
   • 使用torch.gt(importance, threshold)得到掩码
   • 对张量进行原地更新（tensor.mul_(mask)），将小于阈值的权重置零，完成张量剪枝。

def fine_grained_prune(tensor: torch.Tensor, sparsity : float) -> torch.Tensor:"""magnitude-based pruning for single tensor:param tensor: torch.(cuda.)Tensor, weight of conv/fc layer:param sparsity: float, pruning sparsitysparsity = #zeros / #elements = 1 - #nonzeros / #elements:return:torch.(cuda.)Tensor, mask for zeros"""# 确保稀疏率在 [0.0, 1.0] 之间sparsity = min(max(0.0, sparsity), 1.0)	# 如果 sparsity 为 1.0，所有元素被置零并返回全零掩码			if sparsity == 1.0:tensor.zero_()return torch.zeros_like(tensor)# 如果 sparsity 为 0.0，返回全一掩码（不进行剪枝）elif sparsity == 0.0:return torch.ones_like(tensor)num_elements = tensor.numel()num_zeros = round(num_elements * sparsity)importance = tensor.abs()threshold = importance.view(-1).kthvalue(num_zeros).valuesmask = torch.gt(importance, threshold)tensor.mul_(mask)return mask

2. 定义一个细粒度剪枝类FineGrainedPruner，用于管理剪枝过程。
   • init方法：
     • 接收一个模型和一个剪枝字典sparsity_dict，字典中指定了每层的剪枝比例。
     • 使用静态方法 prune 对模型进行剪枝并保存掩码
   • prune方法：遍历模型的参数，对每个参数应用fine_grained_prune函数，生成并存储掩码，且进行剪枝。
   • apply方法：重新应用剪枝掩码。这通常是为了在训练期间（例如，在每个训练步骤之后）确保剪枝效果仍然有效。

class FineGrainedPruner:def __init__(self, model, sparsity_dict):self.masks = FineGrainedPruner.prune(model, sparsity_dict)@torch.no_grad()def apply(self, model):# 使用保存的掩码对模型的参数进行掩码操作，进行剪枝for name, param in model.named_parameters():if name in self.masks:param *= self.masks[name]@staticmethod@torch.no_grad()def prune(model, sparsity_dict):masks = dict()for name, param in model.named_parameters():if param.dim() > 1: # we only prune conv and fc weightsmasks[name] = fine_grained_prune(param, sparsity_dict[name])return masks

# 3. 设置剪枝字典
sparsity_dict = {'conv1.weight': 0.85,'conv2.weight': 0.8,'fc1.weight': 0.75,'fc2.weight': 0.7,'fc3.weight': 0.8,
}# 4. 应用剪枝
pruner = FineGrainedPruner(model, sparsity_dict)
print(f'After pruning with sparsity dictionary')# 5. 打印模型各层剪枝后的稀疏度
print(f'The sparsity of each layer becomes')
for name, param in model.named_parameters():if name in sparsity_dict:print(f'  {name}: {get_sparsity(param):.2f}')# 6. 测试模型大小和mnist数据集上的准确率
sparse_model_size = get_model_size(model, count_nonzero_only=True)
print(f"Sparse model has size={sparse_model_size / MiB:.2f} MiB = {sparse_model_size / origin_model_size * 100:.2f}% of orgin model size")
sparse_model_accuracy = evaluate(model, test_loader)
print(f"Sparse model has accuracy={sparse_model_accuracy:.2f}% before fintuning")

The sparsity of each layer becomesconv1.weight: 0.85conv2.weight: 0.80fc1.weight: 0.75fc2.weight: 0.70fc3.weight: 0.80
Sparse model has size=0.04 MiB = 26.13% of orgin model size
eval:   0%|          | 0/157 [00:00<?, ?it/s]
Sparse model has accuracy=66.46% before fintuning

# 7. 画出剪枝后模型权重分布图和稀疏度直方图
plot_weight_distribution(model, count_nonzero_only=True)
plot_num_parameters_distribution(model)

4.3.3 对剪枝后的模型进行微调

• 使用 SGD 优化器和交叉熵损失函数对剪枝后的模型进行微调训练
• 即使训练过程中模型可能尝试恢复稀疏权重，所以需要在每个训练迭代结束后，调用 pruner.apply(model) 保持模型的稀疏性
• 每个训练周期后评估模型，并在模型达到新的最佳精度时保存模型的权重。

# 1. 设置训练参数
lr = 0.01
momentum = 0.5
num_finetune_epochs = 5
# 设置SGD优化器和交叉熵损失函数
optimizer = torch.optim.SGD(model.parameters(),  lr=lr, momentum=momentum)  	
criterion = nn.CrossEntropyLoss()  												best_sparse_model_checkpoint = dict()
best_accuracy = 0
print(f'Finetuning Fine-grained Pruned Sparse Model')# 2. 微调模型
for epoch in range(num_finetune_epochs):train(model, train_loader, criterion, optimizer,callbacks=[lambda: pruner.apply(model)])accuracy = evaluate(model, test_loader)is_best = accuracy > best_accuracyif is_best:best_sparse_model_checkpoint['state_dict'] = copy.deepcopy(model.state_dict())best_accuracy = accuracyprint(f'    Epoch {epoch+1} Accuracy {accuracy:.2f}% / Best Accuracy: {best_accuracy:.2f}%')

Finetuning Fine-grained Pruned Sparse Model
train:   0%|          | 0/938 [00:00<?, ?it/s]
eval:   0%|          | 0/157 [00:00<?, ?it/s]Epoch 1 Accuracy 97.48% / Best Accuracy: 97.48%
train:   0%|          | 0/938 [00:00<?, ?it/s]
eval:   0%|          | 0/157 [00:00<?, ?it/s]Epoch 2 Accuracy 97.52% / Best Accuracy: 97.52%
train:   0%|          | 0/938 [00:00<?, ?it/s]
eval:   0%|          | 0/157 [00:00<?, ?it/s]Epoch 3 Accuracy 97.79% / Best Accuracy: 97.79%
train:   0%|          | 0/938 [00:00<?, ?it/s]
eval:   0%|          | 0/157 [00:00<?, ?it/s]Epoch 4 Accuracy 97.79% / Best Accuracy: 97.79%
train:   0%|          | 0/938 [00:00<?, ?it/s]
eval:   0%|          | 0/157 [00:00<?, ?it/s]Epoch 5 Accuracy 97.95% / Best Accuracy: 97.95%

# 3. 评估微调后的最优模型
model.load_state_dict(best_sparse_model_checkpoint['state_dict'])
sparse_model_size = get_model_size(model, count_nonzero_only=True)
print(f"Sparse model has size={sparse_model_size / MiB:.2f} MiB = {sparse_model_size / origin_model_size * 100:.2f}% of dense model size")
sparse_model_accuracy = evaluate(model, test_loader)
print(f"Sparse model has accuracy={sparse_model_accuracy:.2f}% after fintuning")

Sparse model has size=0.04 MiB = 26.13% of dense model size
eval:   0%|          | 0/157 [00:00<?, ?it/s]
Sparse model has accuracy=97.95% after fintuning

# 4. 绘制微调后模型权重分布图和稀疏度直方图
plot_weight_distribution(model)
plot_num_parameters_distribution(model)

4.3.3.1 PyTorch保存模型的两种方式

  在 Python 中，当你将一个对象赋值给另一个变量时，默认情况下它们共享同一个内存地址。如果你只是做了一个简单的赋值，例如：

best_sparse_model_checkpoint['state_dict'] = model.state_dict()

  这实际上存储的是 model.state_dict() 的引用（内存地址）。因此，如果在后续的训练或修改过程中模型参数发生变化，best_sparse_model_checkpoint 中的 state_dict 也会跟着改变，因为它们指向的是同一个对象。
  copy.deepcopy() 会创建一个独立的深拷贝，也就是将对象的内容复制到一个新的内存空间里。这样一来，即使后续修改了 model 的参数，已经存储的 state_dict 不会受到影响。

  你也可以直接使用 torch.save 保存最优模型。使用 torch.save 时，模型的 state_dict 会被序列化并存储在磁盘上。保存完成后，即使模型参数在内存中被修改，磁盘上的文件也不会受到影响。

# 保存最佳模型
if is_best:torch.save(model.state_dict(), 'best_model.pth')# 恢复模型
model.load_state_dict(torch.load('best_model.pth'))

维度	torch.save	copy.deepcopy
保存位置	磁盘	内存
存取速度	磁盘 I/O 相对较慢，频繁保存和加载模型会影响训练效率。	快速（内存操作）
内存占用	不占用额外内存，适合大规模模型或内存有限的环境	占用额外内存
持久性	持久化存储，程序结束后仍可使用	非持久化，程序结束后数据丢失
适用场景	训练后使用最优模型进行推理或分发时 训练过程中定期保存模型的快照，以备训练中断时恢复	训练过程中的临时保存与恢复，适合在训练中动态保存最佳状态
实现复杂性	需要指定保存路径及文件	简单，直接复制内存中的模型

4.3.3.2 回调函数

  callbacks是深度学习训练中的一种设计模式，用于在训练过程中（如每个epoch或batch结束后）插入自定义的代码逻辑，增强训练的可扩展性和灵活性，而不用修改 train 函数本身。例如，使用callbacks可以实现以下功能：

• 保存检查点：在模型达到某个性能阈值时保存模型。
• 早停：当验证集的性能不再提高时停止训练。
• 日志记录：在训练期间记录损失和精度变化。
• 动态调整学习率：根据当前训练的表现动态调整学习率。
• 模型修剪或剪枝：保持模型稀疏性或定期应用剪枝掩码，如示例中pruner.apply(model)。

  在上述代码中，callbacks=[lambda: pruner.apply(model)]表示在每个train()调用结束时执行pruner.apply(model)，确保模型在每次训练循环后保持稀疏性。

4.3.4 对比剪枝前后的模型指标

  下面测试模型的延迟（latency）、计算量（以乘累加操作数，即MACs表示）和模型大小（以参数数量表示）。

# 1. 测量模型延迟
@torch.no_grad()
def measure_latency(model, dummy_input, n_warmup=20, n_test=100):model.eval()# 预热阶段，避免测量结果受到模型权重初始化的影响for _ in range(n_warmup):_ = model(dummy_input)# 真实测试阶段，计算模型的平均延迟t1 = time.time()for _ in range(n_test):_ = model(dummy_input)t2 = time.time()return (t2 - t1) / n_test  # 创建格式化模板，左对齐，宽度为15个字符（不到15填充空格）
table_template = "{:<15} {:<15} {:<15} {:<15}"
print (table_template.format('', 'Original','Pruned','Reduction Ratio'))dummy_input = torch.randn(64, 1, 28, 28)pruned_latency = measure_latency(model, dummy_input)
original_latency = measure_latency(origin_model, dummy_input)
print(table_template.format('Latency (ms)',round(original_latency * 1000, 1),round(pruned_latency * 1000, 1),round(original_latency / pruned_latency, 1)))# 2. 测量模型MACs
original_macs = get_model_macs(origin_model, dummy_input)
pruned_macs = get_model_macs(model, dummy_input)
print(table_template.format('MACs (M)',round(original_macs / 1e6),round(pruned_macs / 1e6),round(original_macs / pruned_macs, 1)))# 3. 测量模型大小（参数量）
original_param = get_num_parameters(origin_model, count_nonzero_only=True).item()
pruned_param = get_num_parameters(model, count_nonzero_only=True).item()
print(table_template.format('Param (M)',round(original_param / 1e6, 2),round(pruned_param / 1e6, 2),round(original_param / pruned_param, 1)))

                Original        Pruned          Reduction Ratio
Latency (ms)    3.8             4.7             0.8            
MACs (M)        18              18              1.0            
Param (M)       0.04            0.01            3.8