Aidge pruning API
=================

.. contents::
   :depth: 2
   :local:

Introduction
------------

Aidge provides pruning functionalities that can be used during learning.

Basic example of training pipeline with pruning in Aidge (for a single epoch):

.. code-block:: python

    # Define model, setup learning...
    model = ...
    dataprovider = ...
    opt = ...

    # Setup pruning
    # 1. Insert Prune nodes
    aidge_pruning.insert_weight_prune_nodes(model)

    # 2. Configure pruning
    pruner = aidge_pruning.IterNonStructPruner()
    pruning_rate = aidge_pruning.constant_pr(0.1)
    pruner.set_pruning_rate_scheduler(pruning_rate)
    pruner.set_data_and_masks(list(aidge_pruning.prune_data_and_masks(model)))

    # 3. Initialize masks
    pruner.update_masks()

    scheduler = aidge_core.SequentialScheduler(model)
    for i, (input, label) in enumerate(tqdm(dataprovider)):
        # Forward pass
        pred = scheduler.forward(data=[input])[0]
        # Reset the gradient
        opt.reset_grad()
        # Compute loss
        loss = aidge_learning.loss.CELoss(pred, label)
        # Compute accuracy (optional)
        acc = aidge_learning.metrics.Accuracy(pred, label, 1)[0]
        # Backward pass
        scheduler.backward()
        # Optimize the parameters
        opt.update()

        # Update masks
        pruner.update_masks()


Components
----------

Pruners
~~~~~~~

The base pruner class is :class:` aidge_pruning.Pruner`:

.. autoclass:: aidge_pruning.Pruner
    :members:
    :inherited-members:

The available pruners are:

.. autoclass:: aidge_pruning.IterNonStructPruner
    :members:

.. autoclass:: aidge_pruning.RandomPruner
    :members:


Pruning rate scheduling
~~~~~~~~~~~~~~~~~~~~~~~

The base pruning rate scheduling class is :class:`aidge_pruning.PRScheduler`:

.. autoclass:: aidge_pruning.PRScheduler
    :members:
    :inherited-members:

The available pruning rate schedulers are:

.. autofunction:: aidge_pruning.constant_pr

.. autofunction:: aidge_pruning.step_pr


Recipes
~~~~~~~

.. autofunction:: aidge_pruning.insert_weight_prune_nodes
    
.. autofunction:: aidge_pruning.prune_data_and_masks


Operators
~~~~~~~~~

.. autofunction:: aidge_pruning.Prune