Aidge learning API
==================

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

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

Aidge provides its own native learning module.
The primary purpose of Aidge Learning is to support Aidge's IR for automated model optimization, notably advanced QAT, across heterogeneous deployment targets:

- Ability to apply automated quantization schemes that are easily verifiable and customizable by users;
- Formalized and fully reproducible support for industry-standard models and their derivatives;
- Focus on custom hardware development, necessitating specific training strategies.


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

.. code-block:: python

    # Define the model (or load an existing ONNX one)
    model = ...
    model.set_backend("cuda")
  
    # Initialize parameters (weights and biases)
    aidge_core.init_producer(model, "Producer-1>(Conv2D|FC)", aidge_core.he_filler)
    aidge_core.init_producer(model, "Producer-2>(Conv2D|FC)", lambda x: aidge_core.constant_filler(x, 0.01))

    # Define the data provider (using either aidge_backend_opencv or torch DataLoader)
    dataprovider = ...

    # Define the optimizer and set the parameters to train
    opt = aidge_learning.SGD(momentum=0.9)
    opt.set_parameters(aidge_core.producers(model))
    
    # Define the learning rate scheduler
    learning_rates = aidge_learning.constant_lr(0.01)
    opt.set_learning_rate_scheduler(learning_rates)

    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(model)
        # 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()


Components
----------

Fillers
~~~~~~~

The :class:`aidge_core.init_producer` method can be used to match in the graph the Producer to be filled with a filler:

Usage example:

.. code-block:: python

    # Initialize the weights
    aidge_core.init_producer(model, "Producer-1>(Conv2D|FC)", aidge_core.he_filler)
    # Initialize the bias
    aidge_core.init_producer(model, "Producer-2>(Conv2D|FC)", lambda x: aidge_core.constant_filler(x, 0.01))

.. note::

   Fillers are part of ``aidge_core`` and not ``aidge_learning``.


The available fillers are:

.. autofunction:: aidge_core.constant_filler

.. autofunction:: aidge_core.normal_filler

.. autofunction:: aidge_core.uniform_filler

.. autofunction:: aidge_core.xavier_uniform_filler

.. autofunction:: aidge_core.xavier_normal_filler

.. autofunction:: aidge_core.he_filler


Losses
~~~~~~

.. autofunction:: aidge_learning.loss.MSE

.. autofunction:: aidge_learning.loss.BCE

.. autofunction:: aidge_learning.loss.CELoss

.. autofunction:: aidge_learning.loss.KD

.. autofunction:: aidge_learning.loss.multiStepCELoss


Optimizers
~~~~~~~~~~

The base optimizer class is :class:` aidge_learning.Optimizer`:

.. autofunction:: aidge_learning.Optimizer

The available optimizers are:

.. autofunction:: aidge_learning.Adam

.. autofunction:: aidge_learning.SGD


Learning rate scheduling
~~~~~~~~~~~~~~~~~~~~~~~~

The base learning rate scheduling class is :class:` aidge_learning.LRScheduler`:

.. autofunction:: aidge_learning.LRScheduler

The available learning rate schedulers are:

.. autofunction:: aidge_learning.constant_lr

.. autofunction:: aidge_learning.step_lr


Metrics
~~~~~~~

.. autofunction:: aidge_learning.metrics.Accuracy