随着更多层和节点大型神经网络的使用，降低其存储和计算成本变得至关重要，尤其是对于某些实时应用程序。 模型压缩可用于解决此问题。
我们很高兴的宣布，基于 NNI 的模型压缩工具发布了试用版本。该版本仍处于试验阶段，根据用户反馈会进行改进。 诚挚邀请您使用、反馈，或有更多贡献。
NNI 提供了易于使用的工具包来帮助用户设计并使用压缩算法。 当前支持基于 PyTorch 的统一接口。 只需要添加几行代码即可压缩模型。 NNI 中也内置了一些流程的模型压缩算法。 用户还可以通过 NNI 强大的自动调参功能来找到最好的压缩后的模型，详见自动模型压缩。 另外，用户还能使用 NNI 的接口，轻松定制新的压缩算法，详见教程。
模型压缩方面的综述可参考：Recent Advances in Efficient Computation of Deep Convolutional Neural Networks。
NNI 提供了几种压缩算法，包括剪枝和量化算法：
剪枝
剪枝算法通过删除冗余权重或层通道来压缩原始网络，从而降低模型复杂性并解决过拟合问题。
| 名称 | 算法简介 |
|---|---|
| Level Pruner | 根据权重的绝对值，来按比例修剪权重。 |
| AGP Pruner | 自动的逐步剪枝（是否剪枝的判断：基于对模型剪枝的效果）参考论文 |
| Lottery Ticket Pruner | "The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks" 提出的剪枝过程。 它会反复修剪模型。 参考论文 |
| FPGM Pruner | Filter Pruning via Geometric Median for Deep Convolutional Neural Networks Acceleration 参考论文 |
| L1Filter Pruner | 在卷积层中具有最小 L1 权重规范的剪枝过滤器（用于 Efficient Convnets 的剪枝过滤器） 参考论文 |
| L2Filter Pruner | 在卷积层中具有最小 L2 权重规范的剪枝过滤器 |
| ActivationAPoZRankFilterPruner | 基于指标 APoZ（平均百分比零）的剪枝过滤器，该指标测量（卷积）图层激活中零的百分比。 参考论文 |
| ActivationMeanRankFilterPruner | 基于计算输出激活最小平均值指标的剪枝过滤器 |
| Slim Pruner | 通过修剪 BN 层中的缩放因子来修剪卷积层中的通道 (Learning Efficient Convolutional Networks through Network Slimming) 参考论文 |
量化
量化算法通过减少表示权重或激活所需的精度位数来压缩原始网络，这可以减少计算和推理时间。
| 名称 | 算法简介 |
|---|---|
| Naive Quantizer | 默认将权重量化为 8 位 |
| QAT Quantizer | 为 Efficient Integer-Arithmetic-Only Inference 量化并训练神经网络。 参考论文 |
| DoReFa Quantizer | DoReFa-Net: 通过低位宽的梯度算法来训练低位宽的卷积神经网络。 参考论文 |
| BNN Quantizer | 二进制神经网络：使用权重和激活限制为 +1 或 -1 的深度神经网络。 参考论文 |
通过简单的示例来展示如何修改 Trial 代码来使用压缩算法。 比如，需要通过 Level Pruner 来将权重剪枝 80%，首先在代码中训练模型前，添加以下内容（完整代码）。
PyTorch 代码
from nni.compression.torch import LevelPruner
config_list = [{ 'sparsity': 0.8, 'op_types': ['default'] }]
pruner = LevelPruner(model, config_list)
pruner.compress()TensorFlow 代码
from nni.compression.tensorflow import LevelPruner
config_list = [{ 'sparsity': 0.8, 'op_types': ['default'] }]
pruner = LevelPruner(tf.get_default_graph(), config_list)
pruner.compress()可使用 nni.compression 中的其它压缩算法。 此算法分别在 nni.compression.torch 和 nni.compression.tensorflow 中实现，支持 PyTorch 和 TensorFlow。 参考 Pruner 和 Quantizer 进一步了解支持的算法。 此外，如果要使用知识蒸馏算法，可参考 KD 示例
函数调用 pruner.compress() 来修改用户定义的模型（在 Tensorflow 中，通过 tf.get_default_graph() 来获得模型，而 PyTorch 中 model 是定义的模型类），并修改模型来插入 mask。 然后运行模型时，这些 mask 即会生效。 mask 可在运行时通过算法来调整。
实例化压缩算法时，会传入 config_list。 配置说明如下。
压缩模型时，用户可能希望指定稀疏率，为不同类型的操作指定不同的比例，排除某些类型的操作，或仅压缩某类操作。 配置规范可用于表达此类需求。 可将其视为一个 Python 的 list 对象，其中每个元素都是一个 dict 对象。
list 中的 dict 会依次被应用，也就是说，如果一个操作出现在两个配置里，后面的 dict 会覆盖前面的配置。
在每个 dict 中，有一些 NNI 压缩算法支持的键值：
- op_types：指定要压缩的操作类型。 'default' 表示使用算法的默认设置。
- op_names：指定需要压缩的操作的名称。 如果没有设置此字段，操作符不会通过名称筛选。
- exclude：默认为 False。 如果此字段为 True，表示要通过类型和名称，将一些操作从压缩中排除。
如果使用量化算法，则需要设置更多键值。 如果使用剪枝算法，则可以忽略这些键值
- quant_types : 字符串列表。
要应用量化的类型，当前支持 "权重"，"输入"，"输出"。 "权重"是指将量化操作应用到 module 的权重参数上。 "输入" 是指对 module 的 forward 方法的输入应用量化操作。 "输出"是指将量化运法应用于模块 forward 方法的输出，在某些论文中，这种方法称为"激活"。
- quant_bits : int 或 dict {str : int}
量化的位宽，键是量化类型，值是量化位宽度，例如：
{
quant_bits: {
'weight': 8,
'output': 4,
},
}
当值为 int 类型时，所有量化类型使用相同的位宽。 例如：
{
quant_bits: 8, # 权重和输出的位宽都为 8 bits
}
dict 还有一些其它键值，由特定的压缩算法所使用。 例如， Level Pruner 需要 sparsity 键，用于指定修剪的量。
配置的简单示例如下：
[
{
'sparsity': 0.8,
'op_types': ['default']
},
{
'sparsity': 0.6,
'op_names': ['op_name1', 'op_name2']
},
{
'exclude': True,
'op_names': ['op_name3']
}
]其表示压缩操作的默认稀疏度为 0.8，但op_name1 和 op_name2 会使用 0.6，且不压缩 op_name3。
一些压缩算法使用 Epoch 来控制压缩进度（如AGP），一些算法需要在每个批处理步骤后执行一些逻辑。 因此提供了另外两个 API。 一个是 update_epoch，可参考下例使用：
TensorFlow 代码
pruner.update_epoch(epoch, sess)PyTorch 代码
pruner.update_epoch(epoch)另一个是 step，可在每个批处理后调用 pruner.step()。 注意，并不是所有的算法都需要这两个 API，对于不需要它们的算法，调用它们不会有影响。
使用下列 API 可轻松将压缩后的模型导出，稀疏模型的 state_dict 会保存在 model.pth 文件中，可通过 torch.load('model.pth') 加载。
pruner.export_model(model_path='model.pth')
mask_dict 和 onnx 格式的剪枝模型（需要指定 input_shape）可这样导出：
pruner.export_model(model_path='model.pth', mask_path='mask.pth', onnx_path='model.onnx', input_shape=[1, 1, 28, 28])为了简化压缩算法的编写，NNI 设计了简单且灵活的接口。 对于 Pruner 和 Quantizer 分别有相应的接口。
要实现新的剪枝算法，根据使用的框架，添加继承于 nni.compression.tensorflow.Pruner 或 nni.compression.torch.Pruner 的类。 然后，根据算法逻辑来重写成员函数。
# TensorFlow 中定制 Pruner。
# PyTorch 的 Pruner，只需将
# nni.compression.tensorflow.Pruner 替换为
# nni.compression.torch.Pruner
class YourPruner(nni.compression.tensorflow.Pruner):
def __init__(self, model, config_list):
"""
建议使用 NNI 定义的规范来配置
"""
super().__init__(model, config_list)
def calc_mask(self, layer, config):
"""
Pruner 需要重载此方法来为权重提供掩码
掩码必须与权重有相同的形状和类型。
将对权重执行 ``mul()`` 操作。
此方法会挂载到模型的 ``forward()`` 方法上。
Parameters
----------
layer: LayerInfo
为 ``layer`` 的权重计算掩码
config: dict
生成权重所需要的掩码
"""
return your_mask
# PyTorch 版本不需要 sess 参数
def update_epoch(self, epoch_num, sess):
pass
# PyTorch 版本不需要 sess 参数
def step(self, sess):
"""
根据需要可基于 bind_model 方法中的模型或权重进行操作
"""
pass对于最简单的算法，只需要重写 calc_mask 函数。 它会接收需要压缩的层以及其压缩配置。 可在此函数中为此权重生成 mask 并返回。 NNI 会应用此 mask。
一些算法根据训练进度来生成 mask，如 Epoch 数量。 Pruner 可使用 update_epoch 来了解训练进度。 应在每个 Epoch 之前调用它。
一些算法可能需要全局的信息来生成 mask，例如模型的所有权重（用于生成统计信息）. 可在 Pruner 类中通过 self.bound_model 来访问权重。 如果需要优化器的信息（如在 Pytorch 中），可重载 __init__ 来接收优化器等参数。 然后 step 可以根据算法来处理或更新信息。 可参考内置算法的源码作为示例。
定制量化算法的接口与剪枝算法类似。 唯一的不同是使用 quantize_weight 替换了 calc_mask。 quantize_weight 直接返回量化后的权重，而不是 mask。这是因为对于量化算法，量化后的权重不能通过应用 mask 来获得。
from nni.compression.torch.compressor import Quantizer
class YourQuantizer(Quantizer):
def __init__(self, model, config_list):
"""
建议使用 NNI 定义的规范来配置
"""
super().__init__(model, config_list)
def quantize_weight(self, weight, config, **kwargs):
"""
quantize 需要重载此方法来为权重提供掩码
此方法挂载于模型的 :meth:`forward`。
Parameters
----------
weight : Tensor
要被量化的权重
config : dict
权重量化的配置
"""
# 此处逻辑生成 `new_weight`
return new_weight
def quantize_output(self, output, config, **kwargs):
"""
重载此方法输出量化
此方法挂载于模型的 `:meth:`forward`。
Parameters
----------
output : Tensor
需要被量化的输出
config : dict
输出量化的配置
"""
# 实现生成 `new_output`
return new_output
def quantize_input(self, *inputs, config, **kwargs):
"""
重载此方法量化输入
此方法挂载于模型的 :meth:`forward`。
Parameters
----------
inputs : Tensor
需要被量化的张量
config : dict
输入量化的配置
"""
# 生成 `new_input` 的代码
return new_input
def update_epoch(self, epoch_num):
pass
def step(self):
"""
Can do some processing based on the model or weights binded
in the func bind_model
"""
pass有时，量化操作必须自定义 backward 函数，例如 Straight-Through Estimator，可如下定制 backward 函数：
from nni.compression.torch.compressor import Quantizer, QuantGrad, QuantType
class ClipGrad(QuantGrad):
@staticmethod
def quant_backward(tensor, grad_output, quant_type):
"""
此方法应被子类重载来提供定制的 backward 函数，
默认实现是 Straight-Through Estimator
Parameters
----------
tensor : Tensor
量化操作的输入
grad_output : Tensor
量化操作输出的梯度
quant_type : QuantType
量化类型，可为 `QuantType.QUANT_INPUT`, `QuantType.QUANT_WEIGHT`, `QuantType.QUANT_OUTPUT`,
可为不同的类型定义不同的行为。
Returns
-------
tensor
量化输入的梯度
"""
# 对于 quant_output 函数，如果张量的绝对值大于 1，则将梯度设置为 0
if quant_type == QuantType.QUANT_OUTPUT:
grad_output[torch.abs(tensor) > 1] = 0
return grad_output
class YourQuantizer(Quantizer):
def __init__(self, model, config_list):
super().__init__(model, config_list)
# 定制 backward 函数来重载默认的 backward 函数
self.quant_grad = ClipGrad如果不定制 QuantGrad，默认的 backward 为 Straight-Through Estimator。 即将推出...
- 在 GitHub 中提交此功能的 Bug；
- 在 GitHub 中提交新功能或改进请求；
- 了解 NNI 中特征工程的更多信息；
- 了解 NNI 中 NAS 的更多信息；
- 了解如何使用 NNI 进行超参数调优；