##########
Pruning v2
##########

The ``enot.pruning_v2`` module is the successor of :ref:`enot.pruning <Pruning>` module.
It solves the same problem,
but has a more powerful engine and will support more operations and types of pruning in the future.

All the basic concepts used in this module are described :ref:`here <Pruning>`.

Module exports the following names:

* :class:`~enot.pruning_v2.PruningIR`
* :class:`~enot.pruning_v2.LossAwareCriterion`

*(label selectors)*

* :class:`~enot.pruning_v2.UniformLabelSelector`
* :class:`~enot.pruning_v2.GlobalScoreLabelSelector`
* :class:`~enot.pruning_v2.BinarySearchLatencyLabelSelector`
* :class:`~enot.pruning_v2.KnapsackLabelSelector`
* :class:`~enot.pruning_v2.SCBOLabelSelector`

The pruning workflow can be described as follows:

#. Create :class:`~enot.pruning_v2.PruningIR` using a model.
#. Create :class:`~enot.pruning_v2.LossAwareCriterion` context manager using PruningIR
   and collect scores inside this context (calibration).
#. Create label selector, for example :class:`~enot.pruning_v2.UniformLabelSelector`
   and select the labels to be pruned using the :func:`~enot.pruning_v2.UniformLabelSelector.select` method.
#. Prune model using :func:`~enot.pruning_v2.PruningIR.prune` method of PruningIR and selected labels.

The pruning code of user-defined model looks like this:

.. code-block:: python

    from enot.pruning_v2 import PruningIR
    from enot.pruning_v2 import LossAwareCriterion
    from enot.pruning_v2 import UniformLabelSelector

    # 1. create PruningIR
    ir = PruningIR(model)

    # 2. create criterion and collect scores (calibration)
    with LossAwareCriterion(ir) as score_collector:
        for sample in train_loader:
            model_output = score_collector(sample)
            loss = loss_function(model_output, sample)
            loss.backward()

    # 3. create label selector and select the labels to be pruned
    label_selector = UniformLabelSelector(pruning_ratio=0.5)
    labels_for_pruning = label_selector.select(ir.snapshot())

    # 4. prune model using selected labels
    ir = ir.prune(labels=labels_for_pruning)
    pruned_model = ir.model


.. autoclass:: enot.pruning_v2.PruningIR
    :members: __init__, prune

.. autoclass:: enot.pruning_v2.LossAwareCriterion
    :members: __init__

.. autoclass:: enot.pruning_v2.UniformLabelSelector
    :members:

.. autoclass:: enot.pruning_v2.GlobalScoreLabelSelector
    :members:

.. autoclass:: enot.pruning_v2.BinarySearchLatencyLabelSelector
    :members:

.. autoclass:: enot.pruning_v2.KnapsackLabelSelector
    :members:

.. autoclass:: enot.pruning_v2.SCBOLabelSelector
    :members:

.. autoclass:: enot.pruning_v2.label_selector.LatencyMeasurementError
    :members: