Saving and loading in Aidge

In this tutorial, we will see how to save and load Aidge models with the aidge_core.save_model() and aidge_core.load_model() methods.

Notebook setup

We will need to import the required modules. This feature’s only dependency is aidge_core but to go further on explanations, we will also import the backend_cpu module.

# Must
import aidge_core

# Optional
import aidge_backend_cpu
import numpy as np

And to test these functions, we will also need a model. For this tutorial we will manually create the model with the following function:

def perceptron_model():
    # model creation
    model = aidge_core.sequential(
        [
            aidge_core.FC(in_channels=32 * 32 * 3, out_channels=512),
            aidge_core.ReLU(),
            aidge_core.FC(in_channels=512, out_channels=256),
            aidge_core.ReLU(),
            aidge_core.FC(in_channels=256, out_channels=128),
            aidge_core.ReLU(),
            aidge_core.FC(in_channels=128, out_channels=10),
        ]
    )
    # backend and datatype definition
    model.set_backend("cpu")
    model.set_datatype(aidge_core.dtype.float32)

    # Fill producers with data
    for node in model.get_ordered_nodes():
        if node.type() == "Producer":
            value = node.get_operator().get_output(0)
            tuple_out = node.output(0)[0]
            if tuple_out[0].type() == "FC" and tuple_out[1] == 1:
                # FC weight
                aidge_core.he_filler(value)
            elif tuple_out[0].type() == "FC" and tuple_out[1] == 2:
                # FC bias
                aidge_core.constant_filler(value, 0.01)
    return model


my_model = perceptron_model()

Now that we have our model we could perform a an inference, we will use it to compare results later.

# Create input tensor
input_tensor = aidge_core.Tensor(np.random.rand(32 * 32 * 3).astype(np.float32))
# Create scheduler
scheduler = aidge_core.SequentialScheduler(my_model)

# Run inference
outputs = scheduler.forward(data=[input_tensor])

Saving and Loading

Saving a model is very straightforward:

# path or name of the file
# indicating the extension is optional
path_to_saved_model = "test_save.aidge"
aidge_core.save_model(my_model, path_to_saved_model)

This will save the model to the file “test_save”. The saved model has the extension “.aidge” but in reality it is compressed folder containing multiple files that describe the saved model.

These files are organized like so:

• graph_info.json : GraphViews information, main graph and micrograph (metaops internal graphs)

• node_info.json : Node information, connections, node attributes, etc

• nodes.pkl : pickled operators

• tensors.safetensors : tensor information saved in safetensors

This saving system is a hybrid of pickle, safetensors and custom save&load logic to faithfully recreate Aidge models.

The only thing that will not be saved will be the intermediary outputs/tensors between nodes.

Now we can load it back:

my_model_loaded = aidge_core.load_model(path_to_saved_model)

And lastly we will perform an inference and compare the results:

# New scheduler
scheduler = aidge_core.SequentialScheduler(my_model_loaded)

# Run inference
loaded_outputs = scheduler.forward(data=[input_tensor])

if not aidge_core.approx_eq(outputs[0], loaded_outputs[0], 0.0, 0.0):
    raise RuntimeError("different outputs")
else:
    print("equal outputs, model saved and loaded correctly")