Skip to content
DnsLister Forum

Where domain hunters compare notes

serde-onnx – strongly-typed ONNX graphs for Rust (export + import, validated at compile/export time)

Hey Rustaceans!

I'm excited to share serde-onnx 0.1 – a framework for working with ONNX models as strongly-typed Rust data structures.

Repo: https://github.com/neurono-ml/serde-onnx Crates.io: https://crates.io/crates/serde-onnx | Docs: https://docs.rs/serde-onnx

The problem

Working with ONNX in Rust today is mostly raw prost + HashMap<String, AttributeProto>. You typo an attribute name, pass Vec<i64> where Vec<f32> was expected, or wire node_1 -> node_3 before node_2 exists – and you only find out at inference time with a corrupt .onnx.

What serde-onnx does

It gives you 3 levels, cleanly decoupled:

  1. Raw bytes <-> Generic graph (proto::decode_model / encode_model) – for inspecting any model.
  2. Generic graph <-> Typed IR (Model/Node/ValueType) – lossless, preserves unknown ops as NodePayload::Raw instead of failing.
  3. Typed ops <-> Your Rust structs (GraphBuilder, ToOnnx, #[derive(OnnxExport)]) – compiler checks attribute names and types for you.

And export_model validates the graph (SSA, defined values, topological order) before writing a single byte, with aggregated errors.

Typed ops so far: Scaler, LinearClassifier/Regressor, TreeEnsemble*, SVM*, Imputer, Normalizer, OneHotEncoder, etc. (ai.onnx.ml opset 4) + core Cast/Reshape/Concat/Gather/Identity.

Derive macro vs manual

For the 80% case, a derive is enough:

use serde_onnx::{OnnxExport, export::export_model}; #[derive(OnnxExport)] #[onnx(op = "Scaler", domain = "ai.onnx.ml")] struct DerivedScaler { #[onnx(attr = "offset")] offset: Option<Vec<f32>>, #[onnx(attr = "scale")] scale: Option<Vec<f32>>, not_exported: String, // ignored } let model = export_model(&DerivedScaler { offset: Some(vec![1.0, 2.0]), scale: Some(vec![0.5]), not_exported: String::new() }, "g")?; 

For pipelines / derived attributes, implement ToOnnx manually with GraphBuilder::emit_op:

impl ToOnnx for MyPipeline { fn to_graph(&self, b: &mut GraphBuilder) -> Result<ValueRef, ExportError> { let input = b.input("X", io.clone())?; let imputed = b.emit_op(&Imputer { .. }, vec![input], 1)?.pop().unwrap(); let scaled = b.emit_op(&Scaler { offset: Some(self.mean.clone()), .. }, vec![imputed], 1)?.pop().unwrap(); b.output(scaled.name().to_string(), io)?; Ok(scaled) } } 

Import is symmetric and never fails on unknown ops:

let decoded = DecodedModel::decode_bytes(&bytes)?; match &decoded.payloads[0] { NodePayload::Scaler(op) => println!("{:?}", op.scale), NodePayload::Raw(node) => println!("forwarding {}", node.op_type), _ => {} } 

No protobuf types leak into the IR.

Status: 0.1.0 – looking for help

This is early but tested (round-trip + proptest suites in tests/, 5 runnable examples). What's NOT typed yet: subgraphs (If/Loop/Scan), FunctionProto, and most of ai.onnx (only 5 ops so far).

I'm looking for:

  • Users / testers: Got a real .onnx file from sklearn/pytorch? Try DecodedModel::decode_file and tell me what breaks. I want to know which ops to prioritize next.
  • Contributors: Especially for adding typed ai.onnx ops and more ergonomic builders. Good first issues are labeled. Architecture is decoupled: src/ir <-> src/ml <-> src/proto.
  • Feedback: Is the ToOnnx / GraphBuilder API ergonomic? Is the derive macro pulling its weight?

If you find it useful, a ⭐ on GitHub helps a lot with visibility. Issues and PRs very welcome!

MIT OR Apache-2.0. Thanks for reading!

Source: r/learnrust · by /u/andreclaudino

Leave a Reply

Your email address will not be published. Required fields are marked *