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:
- Raw bytes <-> Generic graph (
proto::decode_model/encode_model) – for inspecting any model. - Generic graph <-> Typed IR (
Model/Node/ValueType) – lossless, preserves unknown ops asNodePayload::Rawinstead of failing. - 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
.onnxfile from sklearn/pytorch? TryDecodedModel::decode_fileand tell me what breaks. I want to know which ops to prioritize next. - Contributors: Especially for adding typed
ai.onnxops and more ergonomic builders. Good first issues are labeled. Architecture is decoupled:src/ir<->src/ml<->src/proto. - Feedback: Is the
ToOnnx/GraphBuilderAPI 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