ONNX models on Udon (experimental)

This chapter is experimental. [Onnx("…")] passes the input and output as arrays and reads the result back as numbers.

The example model

Model file	gesture.onnx, in the same folder as Judge.cs
Input	pose — 12 float32 elements (shape [1,12])
Output	scores — 4 float32 elements (shape [1,4])
The generated file	Judge.Gesture.Model.cs
How to bake one in	As in How to bake one in

The script carrying the attribute

using Tsukimi;
using UnityEngine;

// The attribute argument is the path to the model file, seen from this script.
// Write two or more attributes and each model gets its own set of names.
[Onnx("gesture.onnx")]
public partial class Judge : TsukimiBehaviour
{
    // The generated file becomes the rest of this class, so partial is required.

    // Where the input values come from. Assign 12 of them in the Inspector.
    public Transform[] joints;

    // A value holding the inputs and outputs together. The generated file declares the type.
    private GestureIO io;

    // The chosen index. Other behaviours can read it.
    public int chosen;

    void Start()
    {
        io = new GestureIO();
        // The arrays are allocated once. After that the same arrays are overwritten.
        // Use the generated constants for the element counts (they change with the model).
        // The array type follows the model's element type (this model is float32, so float[]).
        io.In0 = new float[SizeIn0];
        io.Out0 = new float[SizeOut0];
    }

    public void Decide()
    {
        // Fill every input. An element left unfilled is computed with its previous value.
        for (int i = 0; i < SizeIn0; i++) io.In0[i] = joints[i].localPosition.y;

        // Run it. Every layer advances before this returns.
        RunGesture(io);

        // The results land in io.Out0. Here the index of the largest component is picked.
        chosen = 0;
        for (int i = 1; i < SizeOut0; i++)
            if (io.Out0[i] > io.Out0[chosen]) chosen = i;
    }
}

The generated file

This is an excerpt, with the weights and the layer bodies left out.

// This file was generated from an ONNX model. Manual edits are lost on the next generation.
// Source model: gesture

using UnityEngine;

// The type holding the inputs and outputs. The original tensor names and shapes stay in comments.
public struct GestureIO
{
    public float[] In0;   // pose [1,12]
    public float[] Out0;   // scores [1,4]
}

public partial class Judge
{
    // The element counts used when allocating the arrays. Fixed at compile time.
    public const int SizeIn0 = 12;   // [1,12]
    public const int SizeOut0 = 4;   // [1,4]

    // It is private, so it cannot be called from outside without filling the inputs.
    private void RunGesture(GestureIO io)
    {
        float[] In0 = io.In0;   // pose
        float[] Out0 = io.Out0;   // scores

        // The trained weights appear as assignments into arrays.
        float[] c0 = new float[12];
        c0[0] = 0.0002085293f;
        c0[1] = 0.71990794f;
        // …(one for every weight)…

        // Arrays held between layers. Layers that pack four at a time become Vector4.
        float[] t5 = new float[12];   // centered [1,12]
        Vector4[] t7 = new Vector4[6];   // g0 [1,24]

        // The layers appear one at a time, labelled with their name and operator in the original model.
        // center: Elementwise  pose mean -> centered
        for (int i = 0; i < 12; i++)
        {
            t5[i] = In0[i] - c0[i];
        }

        // …(one for every layer)…

        // out: Copy  g2 -> scores
        for (int i = 0; i < 4; i++)
        {
            Out0[i] = t11[i];
        }
    }
}

The generated names

Generated names	Description
GestureIO	A type holding the inputs and outputs. The contents are arrays of the model’s element type (Element types), named In0, Out0, and so on
SizeIn0 / SizeOut0	The element count of each array. The arrays are allocated once in Start
RunGesture(io)	Runs on the inputs in that value and writes the outputs back into the same value
In0 / In1 …	The number follows the order of the model’s inputs. The outputs follow the same rule

Where it can be called	The method that runs cannot be called from outside. This prevents running it without assigning the inputs
Cost	The computation advances one step at a time on Udon. Measure it with the profiler before running it every frame
Where the weights are kept	The weights are constants in the code. The model file does not have to ship with the world
Outputs that no node computes	A model with an output that no node computes (a constant stored in the model, or an input listed as an output as well) is an error. Read that constant, or the input array, in your own code

Element types

Element type in the model	Input and output arrays
float32	float[]
float64	double[]
uint8 / int8	byte[] / sbyte[]
uint16 / int16	ushort[] / short[]
uint32 / int32	uint[] / int[]
uint64 / int64	ulong[] / long[]
bool	bool[]
float16, bfloat16, float8, 4-bit integers, complex numbers, strings	Udon has no such type, so a single input or weight of that type gives a compile error that names the type (convert a float16 or bfloat16 model to float32 before exporting it)

The arithmetic is done in that type, and integers wrap around on overflow (uint8 0 - 1 is 255; the types each operator accepts are in ONNX model operators (experimental)).