Skip to main content

Export a graph to ONNX

ModelGraph::export_onnx writes a graph as a runnable ONNX model file, and io::OnnxModel::from_graph builds the same model in memory without writing anything. The model keeps the graph's input and output names, dtypes and dimensions. A dynamic dimension keeps its name, so two inputs that share one still share it when the model is read back, and the graph's weights become the model's initializers. io::OnnxModel::open and compile() read the file back into a graph that returns the same values. The calls are available from C++ and Python, under the same names.

The graph to export​

The examples on this page export one elementwise model, y = Relu(x * w + b), over an input x of shape [batch, 37] whose first dimension is dynamic and named batch. The model captures w and b, two vectors of 37 values, so the traced graph holds them as its two weights. input(batch) builds an input whose row i holds i + 1 in every column, and row_sums prints the sum of each output row, so two graphs that return the same values print the same line. Every value on this page is a multiple of 1/8, which float32 holds exactly. print_specs prints an input or an output with a dynamic dimension under its name, and the other helpers name a temporary directory for the files a program writes, list its files and print a refusal. In Python the programs import clika_runtime as crt, and ModelGraph and Transform from clika_runtime.graph. There rows(batch) builds the input, slots prints the (name, dtype, dims) tuples that inputs() and outputs() return, with a dynamic dimension as -1, and a value's spec().dim_names holds the dimension names.

A trace returns the graph as built, with every operator as written and nothing optimized or finalized, and export_onnx writes it in that state. This program prints the graph's input and output, then runs it once after finalize(). Each program on this page is complete and runs on its own. From here on, a block shows the part of its program that follows the opening lines the first block shows (the includes or imports, the model and the helpers).

get_the_graph.cpp
#include <algorithm>
#include <cstdint>
#include <cstdio>
#include <filesystem>
#include <string>
#include <vector>

#include <ClikaRT/clika_rt.h>

using ClikaRT::DataType;
using ClikaRT::Error;
using ClikaRT::Tensor;
using ClikaRT::graph::ModelGraph;
using ClikaRT::graph::OnnxExportOptions;
using ClikaRT::io::OnnxModel;
using ClikaRT::spec::TensorSpec;
namespace fs = std::filesystem;
namespace ops = ClikaRT::ops;

namespace {

constexpr std::int64_t kWidth = 37;   // the length of every row

// y = Relu(x * w + b) over x [batch, 37], whose first dim is dynamic and named batch. The model captures
// w and b, two [37] vectors, so the graph holds them as its two weights: w[j] = (j - 18) / 8, from -2.25
// to 2.25, and b = 0.5. A trace returns the graph as built: every operator as written, nothing optimized
// or finalized.
ModelGraph trace_model() {
    const Tensor w = ops::div(ops::sub(ops::arange(0, kWidth, 1, DataType::Float32), 18), 8);
    const Tensor b = Tensor::full({kWidth}, 0.5, DataType::Float32);
    const auto model = [w, b](const std::vector<Tensor>& inputs) -> std::vector<Tensor> {
        return {ops::relu(ops::add(ops::mul(inputs[0], w), b))};
    };
    const std::vector<TensorSpec> signature = {
        TensorSpec("x", DataType::Float32, {TensorSpec::kDynamicDim, kWidth}, false, {"batch", ""})};
    const std::vector<std::string> outputs = {"y"};
    return ClikaRT::graph::trace(model, signature, "elementwise", outputs);
}

// x [batch, 37], whose row i holds i + 1 in every column.
Tensor input(std::int64_t batch) { return ops::cumsum(Tensor::ones({batch, kWidth}, DataType::Float32), 0); }

// The sum of each row of y, separated by spaces.
std::string row_sums(const Tensor& y) {
    std::string out;
    for (const float value : ops::sum(y, {1}).item_as_vec<float>()) {
        char text[32];
        std::snprintf(text, sizeof(text), "%g", static_cast<double>(value));
        out += (out.empty() ? "" : " ") + std::string(text);
    }
    return out;
}

// Each spec on its own line after `label`, as "<name> <dtype> [<dims>]", a dynamic dim written as its name.
void print_specs(const char* label, const std::vector<TensorSpec>& specs) {
    for (const TensorSpec& spec : specs) {
        std::string dims;
        for (std::size_t d = 0; d < spec.dims.size(); ++d) {
            const bool named = d < spec.dim_names.size() && !spec.dim_names[d].empty();
            dims += (dims.empty() ? "" : ", ") + (named ? spec.dim_names[d] : std::to_string(spec.dims[d]));
        }
        std::printf("%s %s %s [%s]\n", label, spec.name.c_str(), ClikaRT::dtype::data_type_name(spec.dtype),
                    dims.c_str());
    }
}

// An empty directory for the files a program writes, under the system's temporary directory.
fs::path fresh_directory(const std::string& name) {
    const fs::path dir = fs::temp_directory_path() / ("export_a_graph_to_onnx_" + name);
    fs::remove_all(dir);
    fs::create_directories(dir);
    return dir;
}

// The names of the files in dir, sorted and separated by spaces, or "none".
std::string files_in(const fs::path& dir) {
    std::vector<std::string> names;
    for (const fs::directory_entry& entry : fs::directory_iterator(dir)) {
        names.push_back(entry.path().filename().string());
    }
    std::sort(names.begin(), names.end());
    std::string out;
    for (const std::string& name : names) out += (out.empty() ? "" : " ") + name;
    return out.empty() ? "none" : out;
}

// A refusal as "<code> | <message>".
void print_refusal(const Error& error) { std::printf("%s | %s\n", error.code_name().c_str(), error.what()); }

}  // namespace

int main() {
    ModelGraph graph = trace_model();
    print_specs("input", graph.inputs());
    print_specs("output", graph.outputs());
    // input x Float32 [batch, 37]
    // output y Float32 [batch, 37]
    graph.finalize();   // a finalized graph runs
    std::printf("%s\n", row_sums(graph.run({input(2)}).front()).c_str());
    // 31.625 52.5
    return 0;
}

Export the graph and read it back​

graph.export_onnx(path) writes the model to path. OnnxModel::open(path) reads the file, and compile() turns it into a ModelGraph again, as built, so finalize() readies it to run (Run an ONNX model covers that side). The read-back graph has the same inputs and outputs, and its dynamic dimension carries the name batch, so it runs at any batch size. At batch 3 and at batch 5 the two graphs return the same sums.

In Python, export_onnx takes the path as a str or any os.PathLike, such as a pathlib.Path, and its options as keyword arguments. OnnxModel.open takes the path as a str.

export_and_read_back.cpp
int main() {
    const fs::path dir = fresh_directory("export_and_read_back");
    const std::string path = (dir / "elementwise.onnx").string();
    ModelGraph graph = trace_model();
    graph.export_onnx(path);   // writes the graph as traced

    // OnnxModel::open reads the file, and compile() makes it a graph again: the same inputs and outputs,
    // with the dynamic dim under its name.
    ModelGraph back = OnnxModel::open(path).compile();
    print_specs("input", back.inputs());
    print_specs("output", back.outputs());
    // input x Float32 [batch, 37]
    // output y Float32 [batch, 37]

    graph.finalize();
    back.finalize();
    for (const std::int64_t batch : {3, 5}) {   // the dynamic dim takes any size
        const Tensor x = input(batch);
        std::printf("%s | %s\n", row_sums(graph.run({x}).front()).c_str(),
                    row_sums(back.run({x}).front()).c_str());
    }
    // 31.625 52.5 73.75 | 31.625 52.5 73.75
    // 31.625 52.5 73.75 95 116.375 | 31.625 52.5 73.75 95 116.375
    fs::remove_all(dir);
    return 0;
}

Build the model in memory​

OnnxModel::from_graph(graph) builds, in memory, the model export_onnx writes, and writes nothing. The result is an OnnxModel like one that open returns: inspect its inputs, outputs and initializers, edit it, optimize() it, compile() it, or save() it when you want the file. It reads the graph's weights only when it is saved or compiled, so it holds no copy of them, and it keeps them alive for as long as it lives. Here the model reports the two weights as its initializers, and the graph it compiles to returns the same sums. In Python, OnnxModel.from_graph is a static method with the same keyword options, and num_initializers is a property.

in_memory.cpp
int main() {
    ModelGraph graph = trace_model();
    // OnnxModel::from_graph builds the model export_onnx writes, in memory: nothing is written.
    OnnxModel model = OnnxModel::from_graph(graph);
    print_specs("input", model.inputs());
    print_specs("output", model.outputs());
    std::printf("%zu initializers\n", model.num_initializers());
    // input x Float32 [batch, 37]
    // output y Float32 [batch, 37]
    // 2 initializers

    // It compiles like any other model, and save() writes it when you want the file.
    ModelGraph back = model.compile();
    graph.finalize();
    back.finalize();
    const Tensor x = input(2);
    std::printf("%s | %s\n", row_sums(graph.run({x}).front()).c_str(), row_sums(back.run({x}).front()).c_str());
    // 31.625 52.5 | 31.625 52.5
    const fs::path dir = fresh_directory("in_memory");
    model.save((dir / "elementwise.onnx").string());
    std::printf("%s\n", files_in(dir).c_str());
    // elementwise.onnx
    fs::remove_all(dir);
    return 0;
}

Choose where the weights go​

OnnxExportOptions::save_as_external_data decides where the weights go when the model is written. true puts them in a <stem>_data file beside the model, false keeps them inside the model file, and unset, the default, keeps them inside until they pass 2 GiB. OnnxModel::open finds the data file beside the model. Here the data file holds the two weights, 296 bytes, and the model read from it returns the same sums. In Python the option is save_as_external_data=True, False or None.

where_the_weights_go.cpp
int main() {
    const fs::path dir = fresh_directory("where_the_weights_go");
    ModelGraph graph = trace_model();

    // save_as_external_data = true writes the weights to a <stem>_data file beside the model.
    OnnxExportOptions beside;
    beside.save_as_external_data = true;
    graph.export_onnx((dir / "beside.onnx").string(), beside);
    // false keeps them in the model file, and so does the default until they pass 2 GiB.
    OnnxExportOptions inside;
    inside.save_as_external_data = false;
    graph.export_onnx((dir / "inside.onnx").string(), inside);
    graph.export_onnx((dir / "default.onnx").string());
    std::printf("%s\n", files_in(dir).c_str());
    // beside.onnx beside_data default.onnx inside.onnx
    std::printf("beside_data: %llu bytes\n", static_cast<unsigned long long>(fs::file_size(dir / "beside_data")));
    // beside_data: 296 bytes: the two [37] float32 weights

    // OnnxModel::open finds the data file beside the model.
    ModelGraph back = OnnxModel::open((dir / "beside.onnx").string()).compile();
    back.finalize();
    std::printf("%s\n", row_sums(back.run({input(2)}).front()).c_str());
    // 31.625 52.5
    fs::remove_all(dir);
    return 0;
}

Set the opset and the other options​

opset_version is the ai.onnx opset the model imports, and 0 takes the library's default. contrib_ops lets the model use the com.microsoft operators, such as GroupQueryAttention and MatMulNBits, and imports that domain only when an operator uses it. optimize simplifies the written model, composing pairs of transposes that undo each other and dropping the nodes no output reads. producer_name is the producer the model records, and empty takes the library's. OnnxModel::from_graph takes the same options, and an operator the chosen opset cannot state is refused, naming the node. In Python the options are the keyword arguments of the same names.

options.cpp
int main() {
    const fs::path dir = fresh_directory("options");
    const std::string path = (dir / "opset17.onnx").string();
    ModelGraph graph = trace_model();

    OnnxExportOptions options;
    options.opset_version = 17;              // the ai.onnx opset the model imports; 0 takes the library's
    options.producer_name = "my-exporter";   // the producer the model records; empty takes the library's
    graph.export_onnx(path, options);
    std::printf("opset %lld\n", static_cast<long long>(OnnxModel::open(path).opset()));
    // opset 17

    // contrib_ops = false keeps the model to the ai.onnx operators, and optimize = false writes it as
    // lowered, without simplifying it.
    OnnxExportOptions plain;
    plain.contrib_ops = false;
    plain.optimize    = false;
    ModelGraph back = OnnxModel::from_graph(graph, plain).compile();
    back.finalize();
    std::printf("%s\n", row_sums(back.run({input(2)}).front()).c_str());
    // 31.625 52.5
    fs::remove_all(dir);
    return 0;
}

A graph that rewrites a buffer in place​

An ONNX model is a function of its inputs, so it cannot state a buffer that a graph changes in place. A trace of a step that writes a captured cache through an in-place operation surfaces the written cache as an extra output, state_out_0, and export_onnx refuses that graph with NOT_IMPLEMENTED_FOR_PARAM, naming the output that carries the update. A graph given a KV cache by attach_kv_cache() is refused the same way. Export the graph before attach_kv_cache(), or trace the step with the cache as an input and its update as an output. That graph exports, and its model reads the cache in and writes the update out. In Python the refusal raises crt.UnsupportedError.

state_buffer.cpp
int main() {
    const fs::path dir = fresh_directory("state_buffer");
    const std::string path = (dir / "step.onnx").string();

    // A step that adds x into a cache it captures, in place. The trace surfaces the written cache as the
    // output state_out_0, and every run rewrites the cache's buffer, which ONNX cannot state.
    const Tensor cache = Tensor::zeros({2, kWidth}, DataType::Float32);
    const auto in_place = [cache](const std::vector<Tensor>& inputs) -> std::vector<Tensor> {
        Tensor written = cache;   // a handle on the live buffer
        ops::add_(written, inputs[0]);
        return {ops::relu(written)};
    };
    const std::vector<TensorSpec> x = {TensorSpec("x", DataType::Float32, {2, kWidth})};
    const std::vector<std::string> y = {"y"};
    ModelGraph writes = ClikaRT::graph::trace(in_place, x, "step", y);
    print_specs("output", writes.outputs());
    // output y Float32 [2, 37]
    // output state_out_0 Float32 [2, 37]
    try {
        writes.export_onnx(path);
    } catch (const Error& error) {
        print_refusal(error);
    }
    // NOT_IMPLEMENTED_FOR_PARAM | export_onnx: the graph rewrites a state buffer (its update is the output
    // 'state_out_0') in place, which ONNX cannot state; export the graph before attach_kv_cache, or trace
    // the step with the cache as an input and its update as an output
    std::printf("files: %s\n", files_in(dir).c_str());
    // files: none

    // The same step with the cache as an input and its update as an output exports.
    const auto explicit_cache = [](const std::vector<Tensor>& inputs) -> std::vector<Tensor> {
        const Tensor updated = ops::add(inputs[1], inputs[0]);   // cache + x
        return {ops::relu(updated), updated};
    };
    const std::vector<TensorSpec> x_and_cache = {TensorSpec("x", DataType::Float32, {2, kWidth}),
                                                 TensorSpec("cache", DataType::Float32, {2, kWidth})};
    const std::vector<std::string> y_and_cache = {"y", "cache_out"};
    ModelGraph step = ClikaRT::graph::trace(explicit_cache, x_and_cache, "step", y_and_cache);
    step.export_onnx(path);
    const OnnxModel model = OnnxModel::open(path);
    print_specs("input", model.inputs());
    print_specs("output", model.outputs());
    // input x Float32 [2, 37]
    // input cache Float32 [2, 37]
    // output y Float32 [2, 37]
    // output cache_out Float32 [2, 37]
    fs::remove_all(dir);
    return 0;
}

What the export refuses​

A refusal raises ClikaRT::Error, whose code_name() names the code, and writes nothing. Besides a buffer written in place, the export refuses:

  • an operator with no ONNX form, naming the node and the operator (NOT_IMPLEMENTED). A custom operator is one by nature, since its kernel is your own code (Write a custom operator builds one);
  • an opset the library does not know (INVALID_ARGUMENT);
  • a call while optimize(), finalize(), attach_kv_cache(), to() or an edit runs on the graph (FAILED_PRECONDITION, naming the call). This program makes one from inside a transform that optimize() runs (Write a transform).

OnnxModel::from_graph refuses the same way.

In Python a refusal raises a subclass of crt.ClikaRTError whose code_name names the code: crt.UnsupportedError for a code that starts with NOT_IMPLEMENTED, and crt.InvalidArgumentError for INVALID_ARGUMENT and FAILED_PRECONDITION. A path of another type raises TypeError. Python has no surface for writing a custom kernel, so the custom operator's refusal is the C++ tab's, and the Python tab shows the others.

refusals.cpp
// A custom operator, y = 2x: its own shape rule and its own kernel, run on the host.
class TimesTwo final : public ClikaRT::nn::Module {
 public:
    static std::shared_ptr<TimesTwo> make() { return std::shared_ptr<TimesTwo>(new TimesTwo()); }

    Tensor forward(const Tensor& x) const { return this->dispatch({&x, 1})[0]; }

    std::vector<ClikaRT::FakeTensor> output_shapes(
        ClikaRT::Span<const ClikaRT::FakeTensor> inputs) const override {
        return {ClikaRT::FakeTensor(inputs[0].shape(), inputs[0].dtype(), inputs[0].stream())};
    }

    void compute(ClikaRT::Span<const Tensor> inputs, ClikaRT::Span<Tensor> outputs) const override {
        const float* x = static_cast<const float*>(inputs[0].const_data_ptr());
        float* y = static_cast<float*>(outputs[0].mutable_data_ptr());
        const std::int64_t n = inputs[0].numel();
        for (std::int64_t i = 0; i < n; ++i) y[i] = 2.0F * x[i];
    }

 private:
    TimesTwo() = default;
};

int main() {
    const fs::path dir = fresh_directory("refusals");
    const std::string path = (dir / "refused.onnx").string();

    // An operator with no ONNX form is refused by name. A custom operator is one by nature: its kernel is
    // your own code.
    const std::shared_ptr<TimesTwo> twice = TimesTwo::make();
    const auto doubled = [twice](const std::vector<Tensor>& inputs) -> std::vector<Tensor> {
        return {twice->forward(inputs[0])};
    };
    const std::vector<TensorSpec> x = {TensorSpec("x", DataType::Float32, {2, kWidth})};
    const std::vector<std::string> y = {"y"};
    ModelGraph custom = ClikaRT::graph::trace(doubled, x, "doubled", y);
    try {
        custom.export_onnx(path);
    } catch (const Error& error) {
        print_refusal(error);
    }
    // NOT_IMPLEMENTED | export_onnx: node 'Custom_1' (Custom) has no ONNX form: it is a user-defined operator

    // An opset the library does not know.
    ModelGraph graph = trace_model();
    OnnxExportOptions unknown;
    unknown.opset_version = 99;
    try {
        graph.export_onnx(path, unknown);
    } catch (const Error& error) {
        print_refusal(error);
    }
    // INVALID_ARGUMENT | no compatible ONNX IR version for the given opset version

    // A call while another call holds the graph: here export_onnx from inside one of optimize()'s
    // transforms.
    const auto export_inside = [&path](ModelGraph& held) -> ClikaRT::Result<void> {
        try {
            held.export_onnx(path);
        } catch (const Error& error) {
            print_refusal(error);
        }
        return {};
    };
    ClikaRT::graph::OptimizeOptions only_this;
    only_this.transforms = std::vector<ClikaRT::graph::Transform>{
        ClikaRT::graph::Transform::from_function("export_inside", export_inside)};
    graph.optimize(only_this);
    // FAILED_PRECONDITION | export_onnx: optimize() is running on this graph; call export_onnx() after it
    // returns

    std::printf("files: %s\n", files_in(dir).c_str());
    // files: none: a refusal writes nothing
    fs::remove_all(dir);
    return 0;
}