onnx-export-rs 0.1.1

Export canonical Rust machine-learning models to ONNX
Documentation

onnx-export-rs

Export inference-only Rust machine-learning model representations to ONNX. The crate separates canonical weights/tree structures from graph exporters, so an adapter for a training library does not need to implement ONNX itself.

The default export path has a minimum supported Rust version of 1.75 and avoids runtime-specific native dependencies. Enabling the optional validate feature raises the requirement to 1.91, determined by its Tract 0.23 validation dependency; the onnxruntime, smartcore, and linfa features likewise pull heavier toolchains.

Status: the canonical API, exporters, Tract validation, and adapters for smartcore linear/logistic/ridge/Lasso/Elastic Net and linfa linear/binary logistic/multiclass logistic and Linfa classification trees are implemented. SmartCore keeps fitted SVM, tree, forest, and XGBoost internals private. The opt-in smartcore-compat feature provides adapters pinned to SmartCore 0.5.5 by reading its serialized representation. Layout changes are rejected as errors, and SVM adapters require the fitting kernel to be supplied explicitly. The same compatibility layer now covers Euclidean k-nearest neighbors, k-means, Gaussian Naive Bayes, and Extra Trees. Public SmartCore adapters cover PCA and truncated SVD transforms. A parallel opt-in linfa-compat feature reads Linfa 0.8.1's serialized state to cover its Tweedie GLM, PLS regression, Gaussian and Multinomial Naive Bayes, and nonlinear-kernel SVM regression and binary classification.

Format compatibility

Models use ONNX IR 8, core opset 13, and (for trees) ai.onnx.ml opset 3. Protobuf bindings are a small prost-based subset of the stable ONNX wire schema, checked into the crate as Rust rather than built with protoc. This keeps builds reproducible and avoids coupling to private tract-onnx bindings.

Exported numeric parameters are narrowed from f64 to f32. This maximizes runtime compatibility but can lose precision for extreme values. Inputs use a symbolic batch first dimension.

Model family Emitted operators Opset Output
Linear / Ridge / Lasso / ElasticNet Gemm 13 [batch, 1]
Binary logistic Gemm, Sigmoid 13 positive-class probability
Multiclass logistic Gemm, Softmax(axis=1) 13 class probabilities
Decision tree / random forest regression TreeEnsembleRegressor ML 3 target scores
Decision tree / random forest classification TreeEnsembleRegressor, ArgMax ML 3 / core 13 zero-based class index
SVM regression / one-class SVM SVMRegressor ML 1 score
SVM classification SVMClassifier ML 1 integer label and class scores
Gradient-boosted trees TreeEnsembleRegressor ML 3 target, probability, or class index
PCA / truncated SVD / affine transforms Gemm 13 transformed features
k-means / nearest centroid arithmetic reductions, ArgMin 13 zero-based cluster index
k-nearest neighbors arithmetic, TopK, Gather, voting 13 target or integer label
Gaussian Naive Bayes arithmetic reductions, ArgMax, Gather 13 integer label
Multinomial / Bernoulli Naive Bayes Gemm, ArgMax, Gather 13 integer label
Categorical Naive Bayes Gather, score accumulation, ArgMax 13 integer label
DBSCAN radius prediction distance reduction, vote matrix, ArgMax 13 cluster label / noise

Linear and logistic graphs are round-trip tested with Tract 0.23.4. Tract can parse ONNX-ML tree ensembles but currently reports its typed translator as unimplemented, so use ONNX Runtime or another ONNX-ML-capable runtime to execute tree, SVM, and gradient-boosting exports. Enable onnxruntime to run the Microsoft ONNX Runtime-backed integration tests for these operators.

Example

use ndarray::array;
use onnx_export_rs::{canonical::LinearModelWeights, exporters::export_linear, save_to_file};

let weights = LinearModelWeights::new(array![2.0, -1.0], 0.5);
let model = export_linear(&weights);
save_to_file(&model, "linear.onnx")?;
# Ok::<(), onnx_export_rs::Error>(())

Binary logistic models use one coefficient row and n_classes = 2. Multiclass models use one row per class. Tree adapters can recursively build RecursiveNode values and pass them through flatten_tree.

Enable validate for Tract-backed validate_export and compare_predictions. Enable smartcore or linfa for the corresponding fitted-model adapters. Smartcore logistic adapters also expose the ordered class labels associated with the exported probability columns. Enable smartcore-compat for decision_tree_regressor, decision_tree_classifier, random_forest_regressor, random_forest_classifier, xgboost_regressor, svm_regressor, binary svm_classifier, extra_trees_regressor, kmeans, Euclidean knn_regressor/knn_classifier, and all four Naive Bayes variants, dbscan, and standard_scaler in adapters::smartcore_compat. The classifier tree/forest adapters return class labels alongside their canonical forest because tree classification exports produce zero-based class indices.

Enable linfa-compat for tweedie_regressor (generalized linear models), pls, gaussian_naive_bayes, multinomial_naive_bayes, and nonlinear-kernel svm_regressor/binary svm_classifier in adapters::linfa_compat, pinned to Linfa 0.8.1. Unlike the SmartCore SVM adapters, these recover the kernel (Gaussian or polynomial) from the serialized state, so no kernel argument is required; linear-kernel SVMs store no support vectors, so use the public adapters::linfa::linear_svm_score for those. The binary svm_classifier takes the integer labels to emit for Linfa's positive (true) and negative (false) classes, which its boolean targets do not otherwise carry.

Scope

  • K-means and Gaussian Naive Bayes use portable core-ONNX subgraphs because ONNX-ML has no dedicated operators for them. Agglomerative clustering has no out-of-sample prediction in SmartCore and therefore has no inference graph.
  • Agglomerative clustering, Linfa's batch DBSCAN/OPTICS outputs, and other training-only algorithms do not expose a fitted out-of-sample inference model and therefore cannot produce a reusable ONNX inference graph.
  • Gradient boosting is supported through the canonical tree representation; perpetual remains preferable when its native exporter already covers the fitted model.

Licensed under MIT.