ai.onnx

ONNX is the common exchange format for trained machine learning models. Train in PyTorch, TensorFlow or scikit-learn, export to .onnx once, and this package runs the model inside a Salam program with no Python and no ONNX Runtime library. LoadOnnx reads the file, Run computes the outputs from your input tensors, and the results are ordinary tensor values.

Use it for inference: scoring, classifying, embedding. To train or fine tune a network in Salam itself, build it with nn and optim. For large language models in the GGUF format, see gguf and llm.

Core ideas

• A Model is the graph from the file: nodes in the order they run, the weights (inits), and the names of the graph's inputs and outputs.
• feed_names lists the inputs you supply, in order; Run takes one tensor for each and returns one tensor per graph output.
• Nothing stops the program. Every failure sets the ok argument to false and puts the reason in model.error.

Limits to know before you export

• Every tensor is f32 inside. Integer and double weights are converted when the file is read, so an int64 above 16777216 loses precision.
• About 65 operators are implemented, enough for feed forward, convolutional and normalization layers (the list is below). Anything else, such as LSTM, Resize or the quantized operators, fails with unsupported op: NAME.
• Conv has no groups (so no depthwise convolutions), no dilations and only symmetric padding. MaxPool and AveragePool have no padding.
• Weights must be inside the .onnx file. Models exported with external data (a separate .data file, used above 2 GB) do not load.
• Input shapes are not checked against the file and the opset version is ignored. A wrong shape shows up as an error from the first operator that cannot use it.

Checking every step

Both ParseOnnx (or LoadOnnx) and Run report success through their last argument, a bool they set, and leave the reason in the model's error field. Check it after each call: a file that is not ONNX, a graph with an operator this package does not implement, or the wrong number of inputs all come back this way instead of stopping the program. The model in this example is built in memory with encoding.protobuf, a single Relu, so the program runs without any file.

import ai.onnx
import ai.tensor
import compress
import encoding.protobuf

func text(b &: compress.ByteBuf, field: int, value: str):
    protobuf.WriteTag(b, field, protobuf.WT_LEN)
    protobuf.WriteLenPrefixedStr(b, value)
end

func sub(b &: compress.ByteBuf, field: int, message: compress.ByteBuf):
    protobuf.WriteTag(b, field, protobuf.WT_LEN)
    protobuf.WriteLenPrefixedBytes(b, message.data, compress.BufLen(message))
end

func relu_model(): compress.ByteBuf:
    mut node := compress.NewByteBuf()
    text(node, 1, "x")
    text(node, 2, "y")
    text(node, 4, "Relu")
    mut input_x := compress.NewByteBuf()
    text(input_x, 1, "x")
    mut output_y := compress.NewByteBuf()
    text(output_y, 1, "y")
    mut graph := compress.NewByteBuf()
    sub(graph, 1, node)
    sub(graph, 11, input_x)
    sub(graph, 12, output_y)
    mut m := compress.NewByteBuf()
    protobuf.WriteTag(m, 1, protobuf.WT_VARINT)
    protobuf.WriteVarint(m, 8 as u64)
    sub(m, 7, graph)
    ret m
end

func main:
    mut ok := false
    junk := "this is not a model"
    _ := onnx.ParseOnnx(junk as void*, junk.len(), ok)
    println "junk parsed:", ok

    mut bytes := relu_model()
    mut model := onnx.ParseOnnx(bytes.data, compress.BufLen(bytes), ok)
    println "model parsed:", ok, "inputs:", model.feed_names().len()

    mut nothing := Vector {} as Vector<tensor.Tensor>
    _ := onnx.Run(model, nothing, ok)
    println "run without inputs:", ok, "|", model.error

    mut feeds := Vector {} as Vector<tensor.Tensor>
    feeds.push(tensor.Arange(4) - tensor.Full(tensor.Shape1(4), 2.0))
    results := onnx.Run(model, feeds, ok)
    println "run with x:", ok, tensor.ToString(results.get(0))
    model.free()
    compress.BufFree(bytes)
end

feed_names lists the inputs you must supply, in order. It leaves out the graph inputs that are really weights stored in the file, which older exporters list as inputs too.

Looking inside a model file

Before wiring a model into a program, look at what it expects. This tool prints the inputs and outputs, how many nodes of each operator the graph has, and how many numbers the weights hold. Run it on any .onnx file:

salam build onnx-guide-inspect.salam
./onnx-guide-inspect model.onnx

import ai.onnx
import os

func main:
    args := os.Args()
    if args.len() < 2:
        println "usage: onnx-guide-inspect model.onnx"
        ret 1
    end
    mut ok := false
    mut model := onnx.LoadOnnx(args.get(1), ok)
    if not ok:
        println "could not read the model:", model.error
        ret 1
    end
    println "IR version", model.ir_version
    println "feeds", model.feed_names(), "outputs", model.outputs
    mut counts := HashMap {} as HashMap<str, int>
    mut order := Vector {} as Vector<str>
    each node in model.nodes:
        if counts.has(node.op):
            counts.put(node.op, counts.get(node.op) + 1)
        else:
            counts.put(node.op, 1)
            order.push(node.op)
        end
    end
    each name in order:
        println "  " + name, counts.get(name)
    end
    mut weights := 0
    each t in model.inits:
        weights += t.numel()
    end
    println model.inits.len(), "weight tensors,", weights, "numbers"
    model.free()
end

Run on the small model that the reference example at the bottom of this page writes, it prints:

IR version 8
feeds ["x"] outputs ["y"]
  Mul 1
  LeakyRelu 1
1 weight tensors, 3 numbers

Check the operator counts against the supported list before you build anything on a model:

Running a trained classifier

Train the model where it is convenient, export it once, and run it from Salam with no Python at run time. With PyTorch, a classifier of the iris flowers with four inputs and three outputs is exported like this:

torch.onnx.export(model, torch.zeros(1, 4), "iris.onnx",
    input_names=["x"], output_names=["logits"],
    dynamic_axes={"x": {0: "batch"}})

The program feeds two flowers at once as a 2x4 tensor, turns the logits into probabilities with tensor.Softmax and takes the most likely class of each row.

import ai.onnx
import ai.tensor

func main:
    mut ok := false
    mut model := onnx.LoadOnnx("iris.onnx", ok)
    if not ok:
        println "could not read iris.onnx:", model.error
        ret 1
    end
    species := ["setosa", "versicolor", "virginica"]
    mut features := Vector {} as Vector<f32>
    each v in [5.1, 3.5, 1.4, 0.2, 6.7, 3.0, 5.2, 2.3]:
        features.push(v as f32)
    end
    mut shape := tensor.Shape2(2, 4)
    mut feeds := Vector {} as Vector<tensor.Tensor>
    feeds.push(tensor.FromArray(shape, features))
    outputs := onnx.Run(model, feeds, ok)
    if not ok:
        println "the model did not run:", model.error
        ret 1
    end
    mut probs := tensor.Softmax(outputs.get(0))
    best := tensor.ArgMaxAxis(probs, 1, false)
    repeat 2 in i:
        println "flower", i, "is", species[best.at(i) as int]
    end
    model.free()
end

The dynamic_axes setting is what lets the batch size change: without it PyTorch may write the batch size of 1 into Reshape and Expand nodes as a constant, and a batch of two then fails inside the graph. The model must be given the same preprocessing it was trained with, here the raw measurements in centimetres.

Types

struct Attr

import ai.onnx
import ai.tensor
import compress
import encoding.protobuf
import archive
import os

func submessage(b &: compress.ByteBuf, field: int, message: compress.ByteBuf):
    protobuf.WriteTag(b, field, protobuf.WT_LEN)
    protobuf.WriteLenPrefixedBytes(b, message.data, compress.BufLen(message))
end

func text(b &: compress.ByteBuf, field: int, value: str):
    protobuf.WriteTag(b, field, protobuf.WT_LEN)
    protobuf.WriteLenPrefixedStr(b, value)
end

func number(b &: compress.ByteBuf, field: int, value: i64):
    protobuf.WriteTag(b, field, protobuf.WT_VARINT)
    protobuf.WriteVarint(b, value as u64)
end

func node(b &: compress.ByteBuf, op: str, input1: str, input2: str, output: str):
    text(b, 1, input1)
    if input2 != "":
        text(b, 1, input2)
    end
    text(b, 2, output)
    text(b, 4, op)
end

func model(m &: compress.ByteBuf):
    mut graph := compress.NewByteBuf()
    mut mul := compress.NewByteBuf()
    node(mul, "Mul", "x", "w", "h")
    submessage(graph, 1, mul)
    mut relu := compress.NewByteBuf()
    node(relu, "LeakyRelu", "h", "", "y")
    mut alpha := compress.NewByteBuf()
    text(alpha, 1, "alpha")
    number(alpha, 20, 1)
    protobuf.WriteTag(alpha, 2, protobuf.WT_FIXED32)
    protobuf.WriteFloat(alpha, 0.1)
    submessage(relu, 5, alpha)
    submessage(graph, 1, relu)
    mut weight := compress.NewByteBuf()
    number(weight, 1, 3)
    number(weight, 2, 1)
    mut payload := compress.NewByteBuf()
    each factor in [2.0, 2.0, 2.0]:
        protobuf.WriteFloat(payload, factor)
    end
    submessage(weight, 4, payload)
    text(weight, 8, "w")
    submessage(graph, 5, weight)
    mut input_x := compress.NewByteBuf()
    text(input_x, 1, "x")
    submessage(graph, 11, input_x)
    mut output_y := compress.NewByteBuf()
    text(output_y, 1, "y")
    submessage(graph, 12, output_y)
    number(m, 1, 8)
    submessage(m, 7, graph)
end

func main:
    mut buffer := compress.NewByteBuf()
    model(buffer)
    mut good := false
    mut m := onnx.ParseOnnx(buffer.data, compress.BufLen(buffer), good)
    println good, m.ok, m.ir_version, m.nodes.len(), m.inits.len(), m.init_names.get(0)
    mut relu := m.nodes.get(1)
    println relu.op, relu.inputs.get(0), relu.outputs.get(0), relu.attr_index("alpha")
    println (relu.attr_float("alpha", 0.0) * 100.0 + 0.5) as int, relu.attr_int("axis", -1)
    attr := relu.attrs.get(0)
    println attr.name, attr.kind
    println m.feed_names().len(), m.inputs.get(0), m.outputs.get(0)

    mut shape := tensor.Shape2(1, 3)
    mut x := tensor.Full(shape, -1.5)
    x.set_at(2, 4.0)
    mut inputs := Vector {} as Vector<tensor.Tensor>
    inputs.push(x)
    mut outputs := onnx.Run(m, inputs, good)
    mut y := outputs.get(0)
    println good, (y.at(0) * 100.0 - 0.5) as int, y.at(2)

    path := os.TempDir() + "/salam_tiny_en.onnx"
    _ := archive.WriteBinaryFileBuf(path, buffer)
    mut m2 := onnx.LoadOnnx(path, good)
    println good, m2.nodes.len()
    m2.free()
    m.free()
    compress.BufFree(buffer)
end

func free()

struct Node

func attr_index(name: str): int

func attr_int(name: str, dflt: i64): i64

func attr_float(name: str, dflt: f64): f64

func free()

struct Model

func feed_names(): Vector<str>

func free()

Functions

func ParseOnnx(p: void*, n: int, ok &: bool): Model

func LoadOnnx(path: str, ok &: bool): Model

func Run(model &: Model, inputs &: Vector<tensor.Tensor>, ok &: bool): Vector<tensor.Tensor>