{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Customize Basic Pruner\n\nUsers can easily customize a basic pruner in NNI. A large number of basic modules have been provided and can be reused.\nFollow the NNI pruning interface, users only need to focus on their creative parts without worrying about other regular modules.\n\nIn this tutorial, we show how to customize a basic pruner.\n\n## Concepts\n\nNNI abstracts the basic pruning process into three steps, collecting data, calculating metrics, allocating sparsity.\nMost pruning algorithms rely on a metric to decide where should be pruned. Using L1 norm pruner as an example,\nthe first step is collecting model weights, the second step is calculating L1 norm for weight per output channel,\nthe third step is ranking L1 norm metric and masking the output channels that have small L1 norm.\n\nIn NNI basic pruner, these three step is implement as ``DataCollector``, ``MetricsCalculator`` and ``SparsityAllocator``.\n\n-   ``DataCollector``: This module take pruner as initialize parameter.\n    It will get the relevant information of the model from the pruner,\n    and sometimes it will also hook the model to get input, output or gradient of a layer or a tensor.\n    It can also patch optimizer if some special steps need to be executed before or after ``optimizer.step()``.\n\n-   ``MetricsCalculator``: This module will take the data collected from the ``DataCollector``,\n    then calculate the metrics. The metric shape is usually reduced from the data shape.\n    The ``dim`` taken by ``MetricsCalculator`` means which dimension will be kept after calculate metrics.\n    i.e., the collected data shape is (10, 20, 30), and the ``dim`` is 1, then the dimension-1 will be kept,\n    the output metrics shape should be (20,).\n\n-   ``SparsityAllocator``: This module take the metrics and generate the masks.\n    Different ``SparsityAllocator`` has different masks generation strategies.\n    A common and simple strategy is sorting the metrics' values and calculating a threshold according to the configured sparsity,\n    mask the positions which metric value smaller than the threshold.\n    The ``dim`` taken by ``SparsityAllocator`` means the metrics are for which dimension, the mask will be expanded to weight shape.\n    i.e., the metric shape is (20,), the corresponding layer weight shape is (20, 40), and the ``dim`` is 0.\n    ``SparsityAllocator`` will first generate a mask with shape (20,), then expand this mask to shape (20, 40).\n\n## Simple Example: Customize a Block-L1NormPruner\n\nNNI already have L1NormPruner, but for the reason of reproducing the paper and reducing user configuration items,\nit only support pruning layer output channels. In this example, we will customize a pruner that supports block granularity for Linear.\n\nNote that you don't need to implement all these three kinds of tools for each time,\nNNI supports many predefined tools, and you can directly use these to customize your own pruner.\nThis is a tutorial so we show how to define all these three kinds of pruning tools.\n\nCustomize the pruning tools used by the pruner at first.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import torch\nfrom nni.algorithms.compression.v2.pytorch.pruning.basic_pruner import BasicPruner\nfrom nni.algorithms.compression.v2.pytorch.pruning.tools import (\n    DataCollector,\n    MetricsCalculator,\n    SparsityAllocator\n)\n\n\n# This data collector collects weight in wrapped module as data.\n# The wrapped module is the module configured in pruner's config_list.\n# This implementation is similar as nni.algorithms.compression.v2.pytorch.pruning.tools.WeightDataCollector\nclass WeightDataCollector(DataCollector):\n    def collect(self):\n        data = {}\n        # get_modules_wrapper will get all the wrapper in the compressor (pruner),\n        # it returns a dict with format {wrapper_name: wrapper},\n        # use wrapper.module to get the wrapped module.\n        for _, wrapper in self.compressor.get_modules_wrapper().items():\n            data[wrapper.name] = wrapper.module.weight.data\n        # return {wrapper_name: weight_data}\n        return data\n\n\nclass BlockNormMetricsCalculator(MetricsCalculator):\n    def __init__(self, block_sparse_size):\n        # Because we will keep all dimension with block granularity, so fix ``dim=None``,\n        # means all dimensions will be kept.\n        super().__init__(dim=None, block_sparse_size=block_sparse_size)\n\n    def calculate_metrics(self, data):\n        data_length = len(self.block_sparse_size)\n        reduce_unfold_dims = list(range(data_length, 2 * data_length))\n\n        metrics = {}\n        for name, t in data.items():\n            # Unfold t as block size, and calculate L1 Norm for each block.\n            for dim, size in enumerate(self.block_sparse_size):\n                t = t.unfold(dim, size, size)\n            metrics[name] = t.norm(dim=reduce_unfold_dims, p=1)\n        # return {wrapper_name: block_metric}\n        return metrics\n\n\n# This implementation is similar as nni.algorithms.compression.v2.pytorch.pruning.tools.NormalSparsityAllocator\nclass BlockSparsityAllocator(SparsityAllocator):\n    def __init__(self, pruner, block_sparse_size):\n        super().__init__(pruner, dim=None, block_sparse_size=block_sparse_size, continuous_mask=True)\n\n    def generate_sparsity(self, metrics):\n        masks = {}\n        for name, wrapper in self.pruner.get_modules_wrapper().items():\n            # wrapper.config['total_sparsity'] can get the configured sparsity ratio for this wrapped module\n            sparsity_rate = wrapper.config['total_sparsity']\n            # get metric for this wrapped module\n            metric = metrics[name]\n            # mask the metric with old mask, if the masked position need never recover,\n            # just keep this is ok if you are new in NNI pruning\n            if self.continuous_mask:\n                metric *= self._compress_mask(wrapper.weight_mask)\n            # convert sparsity ratio to prune number\n            prune_num = int(sparsity_rate * metric.numel())\n            # calculate the metric threshold\n            threshold = torch.topk(metric.view(-1), prune_num, largest=False)[0].max()\n            # generate mask, keep the metric positions that metric values greater than the threshold\n            mask = torch.gt(metric, threshold).type_as(metric)\n            # expand the mask to weight size, if the block is masked, this block will be filled with zeros,\n            # otherwise filled with ones\n            masks[name] = self._expand_mask(name, mask)\n            # merge the new mask with old mask, if the masked position need never recover,\n            # just keep this is ok if you are new in NNI pruning\n            if self.continuous_mask:\n                masks[name]['weight'] *= wrapper.weight_mask\n        return masks"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Customize the pruner.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "class BlockL1NormPruner(BasicPruner):\n    def __init__(self, model, config_list, block_sparse_size):\n        self.block_sparse_size = block_sparse_size\n        super().__init__(model, config_list)\n\n    # Implement reset_tools is enough for this pruner.\n    def reset_tools(self):\n        if self.data_collector is None:\n            self.data_collector = WeightDataCollector(self)\n        else:\n            self.data_collector.reset()\n        if self.metrics_calculator is None:\n            self.metrics_calculator = BlockNormMetricsCalculator(self.block_sparse_size)\n        if self.sparsity_allocator is None:\n            self.sparsity_allocator = BlockSparsityAllocator(self, self.block_sparse_size)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Try this pruner.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Define a simple model.\nclass TestModel(torch.nn.Module):\n    def __init__(self) -> None:\n        super().__init__()\n        self.fc1 = torch.nn.Linear(4, 8)\n        self.fc2 = torch.nn.Linear(8, 4)\n\n    def forward(self, x):\n        return self.fc2(self.fc1(x))\n\nmodel = TestModel()\nconfig_list = [{'op_types': ['Linear'], 'total_sparsity': 0.5}]\n# use 2x2 block\n_, masks = BlockL1NormPruner(model, config_list, [2, 2]).compress()\n\n# show the generated masks\nprint('fc1 masks:\\n', masks['fc1']['weight'])\nprint('fc2 masks:\\n', masks['fc2']['weight'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "This time we successfully define a new pruner with pruning block granularity!\nNote that we don't put validation logic in this example, like ``_validate_config_before_canonical``,\nbut for a robust implementation, we suggest you involve the validation logic.\n\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.8"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}