Training a simple Neural Network

⚠️ ATTENTION: Please compile the packages from source to run this tutorial, as it cannot be currently executed on Binder due to an issue with the aidge-learning PyPI package version 0.5.0. (Last update: 05/12/2025).

This tutorial introduces the basics of learning with the framework Aidge.

What you will learn:

1. creating an SNN model using Aidge API

2. create and import a dataset

3. train the model

Install requirements

Ensure that the Aidge modules along with other required modules such as torch and torchvision, are properly installed in the current environment. If it is the case, the following setup steps can be skipped.

Note: When running this notebook on Binder, all required components are pre-installed.

%pip install aidge-core \
    aidge-backend-cpu \
    aidge-learning

%pip install torch torchvision

Next, we start by importing every module we need:

import aidge_core
import aidge_backend_cpu
import aidge_learning

import numpy as np

# Required to load MNIST dataset
import torchvision
import torchvision.transforms as transforms

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

import torch.nn as nn
import torch

# Changing Log level to avoid too much logs
aidge_core.Log.set_console_level(aidge_core.Level.Fatal)

Creating Aidge model

In this example, we will create a simple perceptron model (One FC (Fully Connected) layer followed by one LIF (Leaky Integrate-And-Fire) layer followed by one FC layer followed by one LIF layer).

For this we cannot use the sequential helper function (as there is currently no way for it to distinguish between the multiple outputs of the Leaky operator), so we create nodes and connect them by hand. We also need to add nodes that will initialize the memory potential of LIF nodes.

First we define the network constants (and in particular the number of timesteps used):

num_steps = 10

# Network topology
in_channels = 28 * 28
hidden_channels = 1000
out_channels = 10

# SNN parameters
beta = 0.95
threshold = 1.0

batch_size = 32

reset_type = aidge_core.leaky_reset.subtraction  # Type of reset for leaky

Then the network itself, and the initialization of the weights and bias.

# Creation of the network
stack = aidge_core.Stack(num_steps, name="stack")
pop = aidge_core.Pop(name="pop")
fc1 = aidge_core.FC(in_channels, hidden_channels, name="fc1")
lif1 = aidge_core.Leaky(num_steps + 1, beta, threshold, reset_type, False, name="lif1")
fc2 = aidge_core.FC(hidden_channels, out_channels, name="fc2")
lif2 = aidge_core.Leaky(num_steps + 1, beta, threshold, reset_type, True, name="lif2")

model = aidge_core.GraphView()

# Connect operators
pop.add_child(fc1, 0, 0)
fc1.add_child(lif1, 0, 0)
lif1.add_child(fc2, 1, 0)  # 1 for spike output
fc2.add_child(lif2, 0, 0)
lif2.add_child(stack, 0, 0)

for name, channels in [("lif1", hidden_channels), ("lif2", out_channels)]:
    memory_init = aidge_core.Tensor()
    memory_init.set_backend("cpu")
    memory_init.to_dtype(aidge_core.dtype.float32)
    memory_init.resize([batch_size, channels])
    memory_init.zeros()
    memory_node = aidge_core.Producer(memory_init, "memory_init_n")

    if name == "lif1":
        memory_node.add_child(lif1, 0, 1)
        memory_node.add_child(lif1, 0, 2)
    else:
        memory_node.add_child(lif2, 0, 1)
        memory_node.add_child(lif2, 0, 2)

    model.add(memory_node)


model.add(fc1)
model.add(lif1)
model.add(fc2)
model.add(lif2)
model.add(stack)
model.add(pop)

model.compile("cpu", aidge_core.dtype.float32)
model.forward_dims()

# Initialize parameters (weights and bias)
for node in model.get_nodes():
    if node.type() == "Producer":
        prod_op = node.get_operator()
        value = prod_op.get_output(0)
        tuple_out = node.output(0)[0]

        # No conv in current network
        if tuple_out[0].type() == "FC" and tuple_out[1] == 1:
            aidge_core.he_filler(value)  # FC weight
        elif tuple_out[0].type() == "FC" and tuple_out[1] == 2:
            aidge_core.constant_filler(value, 0.01)  # FC bias
        else:
            pass

Following the creation of our model, we setup the database object:

def create_mnist_loader():
    dtype = torch.float
    device = torch.device("cpu")
    data_path = "/tmp/data"

    transform = transforms.Compose(
        [
            transforms.Resize((28, 28)),
            transforms.Grayscale(),
            transforms.ToTensor(),
            transforms.Normalize((0,), (1,)),
        ]
    )

    mnist_train = datasets.MNIST(
        data_path, train=True, download=True, transform=transform
    )
    mnist_test = datasets.MNIST(
        data_path, train=False, download=True, transform=transform
    )

    # Create Dataloaders
    train_loader = DataLoader(mnist_train, batch_size, shuffle=True, drop_last=True)
    test_loader = DataLoader(mnist_test, batch_size, shuffle=True, drop_last=True)

    return train_loader, test_loader

Finally, we need to create the learning objects:

opt = aidge_learning.Adam()
opt.set_learning_rate_scheduler(aidge_learning.constant_lr(0.0005))

opt.set_parameters(
    (
        [
            fc1.get_operator().get_input(1),  # FC1 weight
            fc2.get_operator().get_input(1),  # FC2 bias
            fc1.get_operator().get_input(2),  # FC2 weight
            fc2.get_operator().get_input(2),  # FC2 bias
        ]
    )
)

And create the scheduler object:

scheduler = aidge_core.SequentialScheduler(model, False)

Next we can create the training loop. For each batch of data, we call forward(). The difference between this and the normal training is that we also need to setup our scheduler so that is saves inputs and outputs of the nodes (scheduler.save_and_restore_inputs), since they are needed for training. After running forward, we simply call scheduler.backward() to perform BPTT and opt.update() to update the parameters.

def test_learning():

    train_loader, test_loader = create_mnist_loader()
    train_batch = iter(train_loader)
    counter = 0

    for data, targets in train_batch:

        aidge_input = (
            data.view(batch_size, -1)
            .unsqueeze(0)
            .repeat(num_steps, 1, 1)
            .cpu()
            .detach()
            .numpy()
        )
        aidge_label = (
            torch.nn.functional.one_hot(targets, out_channels)
            .numpy()
            .astype(np.float32)
        )

        input_tensor = aidge_core.Tensor(aidge_input)
        input_tensor.to_dtype(aidge_core.dtype.float32)

        for node in model.get_nodes():
            if node.type() == "Pop":
                node.get_operator().associate_input(0, input_tensor)

        scheduler.reset_scheduling()
        scheduler.generate_scheduling()
        scheduler.save_and_restore_inputs(model)
        scheduler.forward()

        output = stack.get_operator().get_output(0)

        target_tensor = aidge_core.Tensor(aidge_label)
        target_tensor.to_dtype(aidge_core.dtype.float32)

        opt.reset_grad()

        loss_value = aidge_learning.loss.multiStepCELoss(
            output, target_tensor, num_steps
        )
        print(f"Loss: {loss_value[0]}", flush=True)

        scheduler.backward()
        opt.update()

        counter += 1
        if counter > 5:
            break

    print("Done.")

Running the learning loop

test_learning()