Custom model creation

Here, we demonstrate the creation of custom ONNX models/pipelines for use within the Nx ecosystem. To export a model from your preferred training platform, please look at the importing model pages.

Note that throughout this section, we heavily rely on the open-source sclblonnx python package; you can find more examples in the sclblonnx git repository.

This article explains how to build ONNX graphs from scratch that encode a pipeline (i.e., an AI model including potential pre- and post-processing) that can be deployed to any edge device running the Nx AI manager.

We cover the following steps:

  1. Basic background regarding the sclblonnx package.

  2. Using sclblonnx to create an ONNX graph from scratch.

If you do not know the ONNX format, we encourage you to read our about ONNX page at this point.

1. Basic background regarding the sclblonnx package

Because at Nx, we use ONNX often, and because our use of ONNX models/pipelines almost always extends (far) beyond simply storing a fitted model in a single environment to use it in that exact same environment later on, we often find ourselves in the situation that we would like to create, inspect, alter, test, or merge existing ONNX graphs. For example, we often add image resizing to an existing vision model such that the resulting ONNX pipeline can be put into production for cameras with different resolutions. However, in our view, the existing onnx.helper API is challenging to use. Thus, internally, we have developed (and are continuously trying to improve) a higher-level API for the manipulation of ONNX graphs. This higher level tooling is openly available in the sclblonnx python package.

The source for the sclblonnx package can be found on git. Easy installation of the package can be done using pip.

In its bare essence, the sclblonnx package provides a number of high-level utility functions to deal with ONNX graphs. We try to use a consistent syntax, which looks as follows:

# Importing the package
import sclblonnx as so

# Assuming we have a graph object g:
g = so.FUNCTION(g, ...)

Thus, we provide a number of functions to operate on a graph (and often alter an existing graph), which results in an updated version of the graph. Common functions are:

  • add_node(g, node): Add a node to an existing graph (and yeah, obviously, you can also delete_node(g, node)).

  • add_input(g, input): Add a new input to an existing graph. You can also delete or change inputs.

  • add_output(g, output): Add a new output to an existing graph.

  • add_constant(g, constant): Add a constant to a graph.

  • clean(g): Clean up the graph; this is important as exported graphs are often bloated or inconsistent.

  • check(g): Check whether the graph is valid, can be run, and can be deployed using the Nx AI Manager (the latter you can turn off)

  • display(g): Visually inspect the graph using Netron.

  • merge(g1, g2, outputs, inputs): Merge two (sub) graphs into a single graph. E.g., add preprocessing to a trained model.

The sclblonnx git repository contains many examples that should help you get started.

2. Using sclblonnx to create an ONNX graph from scratch.

Here we provide the syntax to use the sclblonnx package to create a super simple ONNX graph to add two scalars.

The code example is the first example of creating an ONNX graph from scratch.

Note that the resulting graph cannot be deployed to the AI manager as it does not operate on an image input and does not adhere to our ONNX requirements.

We start by creating an empty graph:

# Use the empty_graph() method to create a named xpb2.GraphProto object:
g = so.empty_graph()

Next, we add the Add node to the graph (you can find the list off all possible nodes, or ONNX operators) here).

# Add a node to the graph.
n1 = so.node('Add', inputs=['x1', 'x2'], outputs=['sum'])
g = so.add_node(g, n1)

By now, we have a graph with a single computational operator called Add. The inputs and output of the add operator are named, but we have not specified their types yet. This is our next step:

# We should explicitly specify the named inputs to the graph -- note that the names determine the graph topology.
# Also, we should specify the data type and dimensions of any input.
# Use so.list_data_types() to see available data types.
g = so.add_input(g, 'x1', "FLOAT", [1])
g = so.add_input(g, 'x2', "FLOAT", [1])

# Similarly, we add the named output with its corresponding type and dimension.
# Note that types will need to "match", as do dimensions. Please see the operator docs for more info.
g = so.add_output(g, 'sum', "FLOAT", [1])

By now we have effectively created a fully functioning ONNX graph: we specified all our operators and the inputs and outputs to the graph (including their types and dimensions).

Next, we provide a few options to check, clean, and inspect the resulting graph:

# so.check() checks the current graph to see if it matches Nx upload criteria for .wasm conversion.
so.check(g)

# Now, a few tricks to sanitize the graph which are always useful.
# so.clean() provides lossless reduction of the graph. If successful cleaned graph is returned.
g = so.clean(g)

# so.display() tries to open the graph using Netron to inspect it. This worsk on most systems if Netron is installed.
# Get Netron at https://github.com/onnx/onnx/blob/master/docs/Operators.md
so.display(g)

If the created graph g passes the so.check() function you can be sure your ONNX graph is proper.

Note: The _sclbl_check argument of the so.check() function can be used to toggle whether or not you would like to check the graph for usage within the Nx ecosystem.

After finalizing and checking the graph, it's easy to test the resulting graph locally using the onnx runtime:

# Now, use the default ONNX runtime to do a test run of the graph.
# Note that the inputs dimensions and types need to match the specification of the graph.
# The outputs returns all the outputs named in the list.
example = {"x1": np.array([1.2]).astype(np.float32), "x2": np.array([2.5]).astype(np.float32)}
result = so.run(g,
                inputs=example,
                outputs=["sum"]
                )
print(result)

Finally, a created graph can easily be stored:

# We can easily store the graph to a file for upload at https://admin.sclbl.nxvms.com/:
so.graph_to_file(g, "onnx/add-scalars.onnx")

After storing a complete graph, you can upload it to the Nx AI Cloud platform by logging into your account at https://admin.sclbl.nxvms.com/ and going to the "CREATE" tab.

After creating your ONNX file, and before uploading it to the Nx platform, please check whether your ONNX file meets all the ONNX requirements imposed by the Nx platform.

Note that the conversion of ONNX to SPMF is one-to-one: i.e., the output produced by the ONNX graph will exactly match the output produced by the AI manager when a model is deployed to any supported edge device.