# 自定义压缩算法

```eval_rst
.. contents::
```

为了简化实现新压缩算法的过程，NNI 设计了简单灵活，同时支持剪枝和量化的接口。 首先会介绍如何自定义新的剪枝算法，然后介绍如何自定义新的量化算法。

**重要说明**，为了更好的理解如何定制新的剪枝、量化算法，应先了解 NNI 中支持各种剪枝算法的框架。 参考[模型压缩框架](https://nni.readthedocs.io/en/latest/Compressor/Framework.html)


## 自定义剪枝算法

要实现新的剪枝算法，需要实现`权重掩码`类，它是 `WeightMasker` 的子类，以及`Pruner` 类，它是 `Pruner` 的子类。

`权重掩码`的实现如下：

```python
class MyMasker(WeightMasker):
    def __init__(self, model, pruner):
        super().__init__(model, pruner)
        # 此处可初始化，如为算法收集计算权重所需要的统计信息。

    def calc_mask(self, sparsity, wrapper, wrapper_idx=None):
        # 根据 wrapper.weight, 和 sparsity, 
        # 及其它信息来计算掩码
        # mask = ...
        return {'weight_mask': mask}
```

参考 NNI 提供的[权重掩码](https://github.com/microsoft/nni/blob/master/src/sdk/pynni/nni/compression/torch/pruning/structured_pruning.py)来实现自己的。

基础的 `Pruner` 如下：

```python
class MyPruner(Pruner):
    def __init__(self, model, config_list, optimizer):
        super().__init__(model, config_list, optimizer)
        self.set_wrappers_attribute("if_calculated", False)
        # 创建权重掩码实例
        self.masker = MyMasker(model, self)

    def calc_mask(self, wrapper, wrapper_idx=None):
        sparsity = wrapper.config['sparsity']
        if wrapper.if_calculated:
            # 如果是一次性剪枝算法，不需要再次剪枝
            return None
        else:
            # 调用掩码函数来实际计算当前层的掩码
            masks = self.masker.calc_mask(sparsity=sparsity, wrapper=wrapper, wrapper_idx=wrapper_idx)
            wrapper.if_calculated = True
            return masks

```

参考 NNI 提供的[Pruner](https://github.com/microsoft/nni/blob/master/src/sdk/pynni/nni/compression/torch/pruning/one_shot.py) 来实现自己的。


***

## 自定义量化算法

要实现新的量化算法，需要继承 `nni.compression.torch.Quantizer`。 然后，根据算法逻辑来重写成员函数。 需要重载的成员函数是 `quantize_weight`。 `quantize_weight` 直接返回量化后的权重，而不是 mask。这是因为对于量化算法，量化后的权重不能通过应用 mask 来获得。

```python
from nni.compression.torch 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):
        """
        根据 bind_model 函数传入的模型或权重进行一些处理
        """
        pass
```

### 定制 backward 函数

有时，量化操作必须自定义 backward 函数，例如 [Straight-Through Estimator](https://stackoverflow.com/questions/38361314/the-concept-of-straight-through-estimator-ste)，可如下定制 backward 函数：

```python
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。 _即将推出_...