#######
Pruning
#######

``enot.pruning`` package contains functional for automatic pruning of user
models.


****************
Supported models
****************

Pruning engine works fine with many models including (but not limited to):

* (from torchvision)
    * classification
        * efficientnet_b0
        * resnet18
        * resnet34
        * resnet50
        * wide_resnet50_2
        * densenet161
        * mobilenet_v2
        * mobilenet_v3_large
    * detection
        * ssd300_vgg16
        * ssdlite320_mobilenet_v3_large
    * segmentation
        * fcn_resnet50
        * fcn_resnet101
        * deeplabv3_resnet50
        * deeplabv3_resnet101
        * deeplabv3_mobilenet_v3_large
        * lraspp_mobilenet_v3_large

* (from Pytorch Image Models (timm))
    * efficientnet_b0
    * resnet18
    * resnet34
    * resnet50
    * wide_resnet50_2
    * densenet161
    * mobilenetv2_100
    * mobilenetv3_large_100


************
Introduction
************

This package features pre-defined or custom pruning procedures for removing
least important filters or neurons. ``enot.pruning`` package currently supports
structured pruning procedure. User can define pruning ratio (percentage of
channels removed) manually or use any of the pre-defined strategies.

The first (and the simplest one) is `equal pruning`_ strategy, which keeps
roughly the same percentage of channels in each prunable layer.

The second (and quite compute intense) is `optimal pruning`_ strategy, which
searches for the optimal pruning ratios for each prunable layer.


A couple definitions to simplify documentation reading experience:

**Pruning ratio** is a fraction of channels to remove (prune) in prunable group
or prunable layer.

**Gate (or channel gate)** is a special ``torch.nn.Module`` which gathers and
saves channel "local importance". "Local importance" means that you can compare
channels by their importance values within single layer, but not between
distinct layers.

**Calibration** procedure estimates channels importance values according to
user model and data.

.. _channel labels:

**Channel label** is the global channel index in a network.

**Pruning config** specifies pruning amount for each prunable layer. See more
:ref:`here <pruning config>`.


***********
Pruning API
***********

Pruning functional accessible to user is divided into three sections:

* `High-level interface`_: the easiest way to use pruning functional;
* `Low-level interface`_: low-level API for pruning, takes moderate time to
  understand and use;
* `Utility functional`_: utility functions for low-level API.


.. _pruning high-level interface:

High-level interface
====================

High-level interface provides two ways to prune model. They work in the same
manner, but have different policies for the amount of pruning for each layer.

.. _equal pruning:

.. autofunction:: enot.pruning.calibrate_and_prune_model_equal

.. _optimal pruning:

.. autofunction:: enot.pruning.calibrate_and_prune_model_optimal


.. _pruning low-level interface:

Low-level interface
===================

The main class for pruning is :class:`.ModelPruningInfo`. It contains necessary
information for model pruning: model execution graph, list of all prunable and
non-prunable channels, all channel importance values, e.t.c.

Model pruning info is obtained through calibration process. Calibration for
pruning is made by :class:`.EnotPruningCalibrator` class. This class's instance
should be used as a context manager. Inside it's context user should compute
losses and perform backward passes through his model on calibration data.

User can also calibrate his model by calling
:func:`.calibrate_model_for_pruning` function and provide all necessary
arguments for proper model execution and loss computation.

Actual model pruning is performed by :func:`.prune_model` function which
requires pruning info object and list of `channel labels`_ to prune.

To select which labels should be pruned, and which should not, we created a
special interface for label selection. Base class for all label selectors is
:class:`.PruningLabelSelector`. This class defines an abstract function
:meth:`.PruningLabelSelector.select` which should select labels to prune and
which subclasses should implement to fit our pipelines based on
:func:`.calibrate_and_prune_model`.

Abstract class :class:`.TopKPruningLabelSelector` is a straightforward way to
select least important channel labels. It utilises :ref:`pruning config` to select
least important channels in each prunable layer. This is done by sorting
channels by their importance values and selecting top-K least important
channels based on the corresponding pruning ratio from the pruning config.

Our label selectors based on this class are the following:

* :class:`.UniformPruningLabelSelector`
* :class:`.OptimalPruningLabelSelector`

They are used internally in high-level interface for pruning described
:ref:`above <pruning high-level interface>`.

**Pruning info class:**

.. autoclass:: enot.pruning.ModelPruningInfo
   :members: summary, n_prunable_groups

**Calibration:**

.. autoclass:: enot.pruning.EnotPruningCalibrator
   :members: __init__, pruning_info

.. autofunction:: enot.pruning.calibrate_model_for_pruning

**Pruning:**

.. autofunction:: enot.pruning.prune_model

The following function combines the functional of
:func:`enot.pruning.calibrate_model_for_pruning` and :func:`enot.pruning.prune_model` functions:

.. autofunction:: enot.pruning.calibrate_and_prune_model

.. _pruning label selectors:

**Label selection:**

.. autoclass:: enot.pruning.PruningLabelSelector
   :members:

.. autoclass:: enot.pruning.TopKPruningLabelSelector
   :members:
   :show-inheritance:

.. autoclass:: enot.pruning.UniformPruningLabelSelector
   :members:
   :show-inheritance:

.. autoclass:: enot.pruning.OptimalPruningLabelSelector
   :members: __init__, select, labels_for_pruning, get_config_for_pruning, get_labels_by_config, pruning_cfg, baseline_latency, minimal_network_latency, n_search_steps


.. _pruning utility functional:

Utility functional
==================

**ModelPruningInfo utilities:**

.. autofunction:: enot.pruning.iterate_over_gate_criteria

.. autofunction:: enot.pruning.get_criteria_label_dict

..
  TODO: Maybe open this function?
  .. autofunction:: enot.pruning.get_total_remaining_importance_value

.. autofunction:: enot.pruning.get_labels_for_uniform_pruning

**Pruning config utilities:**

.. autofunction:: enot.pruning.check_pruning_config

.. autofunction:: enot.pruning.get_least_important_labels_by_config

**OptimalPruningLabelSelector utilities:**

.. autofunction:: enot.pruning.default_min_channels_to_leave

.. autofunction:: enot.pruning.default_channel_search_step