compute-onnx

ONNX Runtime based compute backend for GoMLX.

It allows GoMLX models to be executed via ONNX Runtime using either CPU or CUDA (NVIDIA GPU).

It supports dynamic shapes and exporting models to .onnx files.

Example Usage

To run the Adult dataset demo with the ONNX backend:

GOMLX_BACKEND=onnx go run -tags=onnx github.com/gomlx/gomlx/examples/adult/demo

Or targeting a specific accelerator (e.g. CUDA):

GOMLX_BACKEND=onnx:cuda go run -tags=onnx github.com/gomlx/gomlx/examples/adult/demo

Backend Options & Configuration

Configuration can be specified in the GOMLX_BACKEND environment variable using onnx:<options> or onnxruntime:<options> (comma-separated).

Accelerator Selection

• cpu: Force CPU execution.

  GOMLX_BACKEND=onnx:cpu

• cuda / gpu: Force CUDA GPU execution (uses ONNX Runtime CUDA Execution Provider via OrtIoBinding).

  GOMLX_BACKEND=onnx:cuda

• Custom Library Path: Specify an explicit path to the ONNX Runtime .so (or .dylib / .dll) shared library file. This explicitly bypasses ONNXRUNTIME_SHARED_LIBRARY_PATH.

  GOMLX_BACKEND=onnx:/path/to/libonnxruntime.so

• empty (default): Automatically detects if an NVIDIA GPU is present via nvidia-smi and defaults to CUDA if available, otherwise falling back to CPU.

  GOMLX_BACKEND=onnx

Save Model To ONNX

This allows one to export GOMLX trained (or fine-tuned) models to ONNX.

See an example in UCI-Adult demo. If you have a pre-trained file in a directory called base:

GOMLX_BACKEND=onnx go run -tags=onnx ./examples/adult/demo/ -checkpoint "base" -save_onnx="/tmp/a.onnx" -vmodule=save_onnx=1

ONNX Runtime Shared Libraries & Auto-Installation

The backend automatically locates or manages the required ONNX Runtime shared library (libonnxruntime.so / onnxruntime.dll):

• Custom Library Path: Set the ONNXRUNTIME_SHARED_LIBRARY_PATH environment variable or pass an explicit library path in the backend configuration (e.g. GOMLX_BACKEND=onnx:/path/to/libonnxruntime.so) to point directly to the shared library binary. Passing an explicit path in the configuration bypasses ONNXRUNTIME_SHARED_LIBRARY_PATH.
• Auto-Installation: If no library path is provided, the backend automatically downloads and extracts prebuilt official ONNX Runtime binaries locally (e.g. ~/.local/lib/onnxruntime/ on Linux).
• Disabling Auto-Installation: Set the environment variable GOMLX_NO_AUTO_INSTALL=1 or call onnxbackend.EnableAutoInstall(false) programmatically to disable automatic downloads (useful for offline environments or container deployments).

Debugging

Saving Failed Models (GOMLX_ONNX_SAVE_ON_FAILURE)

If graph compilation or session creation fails in ONNX Runtime, setting the GOMLX_ONNX_SAVE_ON_FAILURE environment variable instructs the backend to automatically save the serialized ONNX model protobuf to the specified file path before returning the compilation error:

GOMLX_ONNX_SAVE_ON_FAILURE="/tmp/failed_model.onnx" go run -tags=onnx ...

This allows you to inspect the invalid graph using onnx_printer or Netron to diagnose the failure.

Inspecting .onnx Files (onnx_printer)

This repository includes a CLI tool in cmd/onnx_printer to inspect and pretty-print .onnx model files in the terminal:

go run github.com/gomlx/compute-onnx/cmd/onnx_printer path/to/model.onnx

It formats input, output, and node tensor shapes using GoMLX shapes.Shape (including named dynamic axes) and prints each graph operation on a single line. Tensor constants and initializers are truncated to 10 elements by default (controlled via -max_items / -n).

Example usage:

# Print model details with a maximum of 5 items for constant values
go run github.com/gomlx/compute-onnx/cmd/onnx_printer -max_items 5 /tmp/model.onnx

# Read from stdin
cat /tmp/model.onnx | go run github.com/gomlx/compute-onnx/cmd/onnx_printer

Tip: For interactive graphical visualization of ONNX models, you can open .onnx model files using Netron.

Logging & Verbosity

• Backend Log Level (log=<level>): Configures ONNX Runtime's internal logging severity level.

  • log=0: Errors only (severity level 3 / ERROR)
  • log=1: Warnings (severity level 2 / WARNING)
  • log=2: Informational (severity level 1 / INFO)
  • log=3: Verbose (severity level 0 / VERBOSE)

  Example:

  GOMLX_BACKEND="onnx:cuda,log=2"

• Execution Timing Log (-vmodule=executable=1): Enables per-step execution timing breakdown printed via klog using humanize.Duration.

  Example:

  GOMLX_BACKEND=onnx:cuda go run -tags=onnx github.com/gomlx/gomlx/examples/adult/demo -vmodule=executable=1

💖 Thanks

• Go
• ONNX
• ONNX Runtime