NNI --模型剪枝

概念

剪枝是压缩神经网络模型的常用技术。 剪枝方法探索模型权重（参数）中的冗余，并尝试删除/修剪冗余和非关键权重。 冗余元素从模型中修剪，其值归零，我们确保它们不参与反向传播过程。

修剪深度学习模型有三种常见做法，NNI 通过在关键修剪阶段工作来支持上述所有修剪实践。

1. 预训练模型 -> 修剪模型 -> 微调修剪后的模型
2. 在训练期间修剪模型（即修剪感知训练） -> 微调修剪后的模型
3. 修剪模型 -> 从头开始训练修剪后的模型

模型剪枝案例

第一步：加载待剪枝的模型
使用一个简单的模型，并在MNIST数据集上进行了预训练。

import torch
import torch.nn.functional as F
from torch.optim import SGD

from nni_assets.compression.mnist_model import TorchModel, trainer, evaluator, device
# 随机数种子
import numpy as np
import random
torch.manual_seed(1)
torch.cuda.manual_seed_all(1)
np.random.seed(1)
random.seed(1)
# define the model
model = TorchModel().to(device)

第二步：待剪枝模型的预训练
训练一个深度学习模型，实际上是在用数据来寻找一个最优的参数组合，这些参数决定了模型在输入数据时如何进行计算。通过训练，我们可以找到一组能够最好地拟合我们的数据的参数。这个过程还可以帮助我们理解每个参数在模型中的重要性。
在模型剪枝的过程中，首先进行预训练的主要原因是我们需要通过训练过程来了解模型中的哪些部分是关键的，哪些部分是冗余的。

# define the optimizer and criterion for pre-training
optimizer = SGD(model.parameters(), 1e-2)
criterion = F.nll_loss

# pre-train and evaluate the model on MNIST dataset
for epoch in range(3):
    trainer(model, optimizer, criterion)
    evaluator(model)

第三步：裁剪的参数设置
通过设置裁剪的参数，可以设置裁剪的稀疏程度、裁剪和不裁剪的op等等。

# 裁剪的参数设置
config_list = [{
    'sparsity_per_layer': 0.5,  # 裁剪的稀疏程度，数值越大，裁剪程度越大
    'op_types': ['Linear', 'Conv2d']    # 裁剪的op
},
{
    'exclude': True,    # 不裁剪的op
    'op_names': ['fc3']     # 一般不裁剪最后一层
}]

第四步：模型裁剪

# 导入NNI库的两种模型修剪器：L1NormPruner（L1范数修剪器）
from nni.compression.pytorch.pruning import L1NormPruner
# 实例化了一个L1NormPruner对象，需要将你要进行剪枝的模型和剪枝的配置列表作为参数。
pruner = L1NormPruner(model, config_list)

# show the wrapped model structure, `PrunerModuleWrapper` have wrapped the layers that configured in the config_list.
# 将打印已经被封装（即进行剪枝的准备）的模型的结构，其中PrunerModuleWrapper封装了在config_list中配置的层。
print(model)

# compress the model and generate the masks
# 使用 L1NormPruner 修剪模型并生成掩码
_, masks = pruner.compress()
# show the masks sparsity
# 遍历所有的mask，并计算它们的稀疏度。稀疏度指的是多少比例的权重被置为零（也就是被修剪）。
for name, mask in masks.items():
    print(name, ' sparsity : ', '{:.2}'.format(mask['weight'].sum() / mask['weight'].numel()))

# need to unwrap the model, if the model is wrapped before speedup
# 解除了模型的封装，返回原始的模型（未经修剪）。在使用ModelSpeedup之前，我们需要解包模型。
pruner._unwrap_model()

# speedup the model, for more information about speedup, please refer :doc:`pruning_speedup`.
# 导入了ModelSpeedup函数，这是一个能够将剪枝操作真正应用到模型上的函数。
from nni.compression.pytorch.speedup import ModelSpeedup

# 加速后模型将变小
ModelSpeedup(model, torch.rand(3, 1, 28, 28).to(device), masks).speedup_model()

第五步：微调压缩模型
如果模型已加速，则需要重新初始化新的优化器以进行微调。 因为加速将用密集的小层取代被遮蔽的大层。

# 微调压缩模型
optimizer = SGD(model.parameters(), 1e-2)
for epoch in range(3):
    trainer(model, optimizer, criterion)
    evaluator(model)