Aidge learning API#

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):

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

About gradient handling #

There is no “autograd” mechanism in Aidge, because the computational graph is explicitly defined, and the gradient computation path is determined from the outset.

Gradients are accumulated by default in the backward implementations of the operators.

Gradients are automatically reset just before each forward pass (it is the first forward hook registered by default for all nodes), except if the node has an attribute grad.no_reset. When producer nodes are added to the optimizer parameters list with aidge_learning.Optimizer.set_parameters, the attribute grad.no_reset is automatically set for these nodes. As a result, their gradients are not reset automatically.

For such nodes, a call to aidge_learning.Optimizer.reset_grad is required to clear the gradients. This allows gradients to be accumulated across multiple forward-backward iterations before being reset.

When performing multiple iterations, make sure to scale the loss accordingly using the scaling argument, as shown in the example below:

# Number of iterations before each optimizer update step
nb_iters = 8

opt.reset_grad()

for i, (input, label) in enumerate(tqdm(dataprovider)):
    pred = scheduler.forward(False, data=[input])[0]
    loss = aidge_learning.loss.CELoss(pred, label, 1 / nb_iters)
    scheduler.backward()

    if (i + 1) % nb_iters == 0 or (i + 1) == len(dataprovider):
        opt.update()
        opt.reset_grad()

    pbar.set_postfix(loss=f"{loss[0]:.4f}")

Data pre-processing #

Image pre-processing and data loading can be performed either directly with Aidge or using external libraries such as Torchvision. Both approaches are supported, allowing you to integrate Aidge into existing training pipelines with minimal changes.

With Aidge #

Aidge provides an OpenCV backend that performs dataset loading and image pre-processing directly inside the data pipeline. Transformations are executed before data reaches the model and can target either the input images or the labels.

For example, the following micro-graph converts integer class labels into one-hot encoded tensors:

label_one_hot = aidge_core.MetaOperatorOp(
    "LabelOneHot",
    aidge_core.sequential(
        [
            aidge_core.OneHot(depth=nb_classes),
            aidge_core.Squeeze(),
            aidge_core.Cast(aidge_core.dtype.float32),
        ]
    ),
)

The dataset is then wrapped in a DataProvider, which is responsible for batching, shuffling, and applying transformations:

train_dataset = aidge_backend_opencv.ILSVRC2012_Directory(
    data_path=args.imagenet_path, train=True
)
train_loader = aidge_core.DataProvider(
    train_dataset,
    backend=args.backend,
    batch_size=args.batch_size,
    shuffle=True,
    drop_last=True,
)
train_loader.set_backend(args.backend, args.dev)

Transformations are attached independently to each dataset channel:

Channel 0: image transformations (data augmentation, color conversion, normalization);
Channel 1: label transformations (one-hot encoding).

train_loader.transforms(
    0,
    aidge_core.sequential(
        [
            aidge_backend_opencv.RandomFlipTransformation(horizontal_flip=0.5),
            aidge_backend_opencv.ColorSpaceTransformation(
                aidge_backend_opencv.colorspace.rgb
            ),
            aidge_backend_opencv.RangeAffineTransformation(
                aidge_backend_opencv.range_affine_operator.minus,
                [123.675, 116.28, 103.53],
                aidge_backend_opencv.range_affine_operator.multiplies,
                [0.0171, 0.0175, 0.0174],
            ),
        ]
    ),
)
train_loader.transforms(
    1,
    label_one_hot.get_micro_graph(),
)

One important point to note is that data transformation graphs are first-class Aidge graphs. They use the exact same formalism as neural network computation graphs, allowing data pre-processing pipelines to be built, manipulated, composed, optimized, and executed just like model graphs.

Visualization of the image transformation pipeline.

Once configured, the training loop directly receives Aidge tensors that are already pre-processed and transferred to the selected backend:

for i, (input, label) in enumerate(pbar):
    pred = scheduler.forward(False, data=[input])[0]

    ...

This approach keeps the entire pre-processing pipeline inside Aidge and avoids additional tensor conversions during training.

With Torchvision #

Existing Torchvision datasets and transforms can also be used without modification. This is particularly useful when migrating an existing PyTorch project to Aidge.

Image transformations are defined using the standard Torchvision API:

from torchvision import datasets, transforms
import torch.utils.data

train_transforms = transforms.Compose(
    [
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
        ),
    ]
)

train_dataset = datasets.ImageFolder(
    f"{args.imagenet_path}/train", transform=train_transforms
)
train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=args.batch_size,
    shuffle=True,
    drop_last=True,
    num_workers=32,
)

Before inference, the PyTorch tensors must be converted to Aidge tensors and transferred to the target backend. Label pre-processing can then be performed using the same Aidge micro-graph as in the previous example.

for i, (input, label) in enumerate(train_loader):
    input = aidge_core.Tensor(input.numpy())
    input.to_backend(args.backend, args.dev)
    label = aidge_core.Tensor(label.numpy())
    label.to_backend(args.backend, args.dev)
    label = label_one_hot.forward(label)[0]

    pred = scheduler.forward(False, data=[input])[0]

    ...

This interoperability allows Aidge to be integrated into existing PyTorch data pipelines while keeping model execution inside Aidge.

Training Monitoring #

With Aidge’s Dashboard #

Aidge provides a built-in training dashboard for monitoring experiments. If preferred, standard visualization tools such as TensorBoard can also be used.

Create a dashboard instance:

monitor = aidge_learning.dashboard.TrainingDashboard(
    path="runs/"
    + args.filename[: -len(".onnx")]
    + datetime.datetime.now().strftime("_%Y%m%d-%H%M%S")
)

Logging scalar values#

Training metrics can be logged with optional visualization settings:

monitor.log(
    "learning_rate",
    opt.learning_rate(),
    step=epoch,
    plot={
        "color": "rgb(192, 75, 75)",
    },
)

monitor.log(
    "train_loss",
    average_loss,
    step=epoch,
    plot={
        "y_begin_at_zero": True,
    },
)

monitor.log("test_accuracy", accuracy, step=epoch)

Result:

Example of train loss during CIFAR-10 training on ResNet-8.

Visualizing input samples#

Image batches can be logged directly for inspection:

for i, (input, label) in enumerate(validation_loader):
    ...
    monitor.log("train_batch", input)
    ...

Result:

Example of train batch sample for CIFAR-10.

Logging confusion matrices#

Metrics such as confusion matrices can also be accumulated and displayed:

confusion_matrix = aidge_learning.metrics.ConfusionMatrix(10)
...

for i, (input, label) in enumerate(validation_loader):
    ...
    pred = scheduler.forward(False, data=[input])[0]
    confusion_matrix.aggregate(label, pred)

monitor.log("confusion_matrix", confusion_matrix, step=epoch)

Result:

Example of confusion matrix on CIFAR-10.

With TensorBoardX #

Aidge can also integrate seamlessly with TensorBoard through the TensorBoardX package.

Create a SummaryWriter:

from tensorboardX import SummaryWriter

board_writer = SummaryWriter(
    log_dir="runs/"
    + args.filename
    + datetime.datetime.now().strftime("_%Y%m%d-%H%M%S")
)

Training metrics can then be logged using the standard TensorBoard API:

board_writer.add_scalar(
    tag="train/learning-rate",
    scalar_value=opt.learning_rate(),
    global_step=epoch,
)

board_writer.add_scalar(
    tag="train/loss",
    scalar_value=np.array(average_loss),
    global_step=epoch,
)

board_writer.add_scalar(
    tag="test/accuracy",
    scalar_value=accuracy,
    global_step=epoch,
)

Aidge learning API#

Introduction #

About gradient handling #

Components #

Fillers #

Losses #

Optimizers #

Learning rate scheduling #

Metrics #

Confusion Matrix#

Data pre-processing #

With Aidge #

With Torchvision #

Training Monitoring #

With Aidge’s Dashboard #

Logging scalar values#

Visualizing input samples#

Logging confusion matrices#

With TensorBoardX #