onnx-export-rs 0.1.1

Export canonical Rust machine-learning models to ONNX
Documentation
use ndarray::{Array1, Array2};

use crate::{Error, Result};

/// Kernel used by an ONNX-ML support-vector machine.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum SvmKernel {
    /// Linear dot-product kernel.
    Linear,
    /// Polynomial kernel.
    Polynomial {
        /// Multiplicative kernel coefficient.
        gamma: f64,
        /// Additive kernel coefficient.
        coef0: f64,
        /// Polynomial degree.
        degree: u32,
    },
    /// Radial-basis-function kernel.
    Rbf {
        /// RBF coefficient.
        gamma: f64,
    },
    /// Sigmoid kernel.
    Sigmoid {
        /// Multiplicative kernel coefficient.
        gamma: f64,
        /// Additive kernel coefficient.
        coef0: f64,
    },
}

impl SvmKernel {
    pub(crate) const fn onnx_name(self) -> &'static [u8] {
        match self {
            Self::Linear => b"LINEAR",
            Self::Polynomial { .. } => b"POLY",
            Self::Rbf { .. } => b"RBF",
            Self::Sigmoid { .. } => b"SIGMOID",
        }
    }

    pub(crate) fn onnx_parameters(self) -> Vec<f32> {
        match self {
            Self::Linear => vec![0.0, 0.0, 0.0],
            Self::Polynomial {
                gamma,
                coef0,
                degree,
            } => {
                vec![gamma as f32, coef0 as f32, degree as f32]
            }
            Self::Rbf { gamma } => vec![gamma as f32, 0.0, 0.0],
            Self::Sigmoid { gamma, coef0 } => vec![gamma as f32, coef0 as f32, 0.0],
        }
    }
}

/// Inference parameters for SVM regression or one-class SVM.
#[derive(Clone, Debug, PartialEq)]
pub struct SvmRegressor {
    /// Shape `[support_count, feature_count]`.
    pub support_vectors: Array2<f64>,
    /// One coefficient per support vector, or one per feature for a compact
    /// linear model with zero support vectors.
    pub coefficients: Array1<f64>,
    /// Bias term added to the kernel score by ONNX Runtime (`rho`).
    pub rho: f64,
    /// Kernel configuration.
    pub kernel: SvmKernel,
    /// Whether this is a one-class SVM.
    pub one_class: bool,
}

impl SvmRegressor {
    /// Constructs a checked SVM regressor representation.
    ///
    /// # Errors
    ///
    /// Returns an error when coefficient and support-vector counts differ or
    /// no input features are present. ONNX's compact linear representation is
    /// accepted when the support-vector matrix has zero rows and one
    /// coefficient is supplied per feature.
    pub fn new(
        support_vectors: Array2<f64>,
        coefficients: Array1<f64>,
        rho: f64,
        kernel: SvmKernel,
        one_class: bool,
    ) -> Result<Self> {
        let compact_linear = support_vectors.nrows() == 0
            && matches!(kernel, SvmKernel::Linear)
            && support_vectors.ncols() == coefficients.len();
        let support_vector_form =
            support_vectors.nrows() != 0 && support_vectors.nrows() == coefficients.len();
        if support_vectors.ncols() == 0 || (!compact_linear && !support_vector_form) {
            return Err(Error::InvalidModel("SVM regressor shape mismatch".into()));
        }
        Ok(Self {
            support_vectors,
            coefficients,
            rho,
            kernel,
            one_class,
        })
    }
}

/// ONNX-ML inference parameters for an integer-labelled SVM classifier.
#[derive(Clone, Debug, PartialEq)]
pub struct SvmClassifier {
    /// Shape `[support_count, feature_count]`.
    pub support_vectors: Array2<f64>,
    /// Flattened ONNX pairwise support-vector coefficients.
    pub coefficients: Array1<f64>,
    /// Pairwise bias terms in ONNX's `rho` convention.
    pub rho: Array1<f64>,
    /// Number of support vectors belonging to each class.
    pub vectors_per_class: Vec<usize>,
    /// Integer class labels in output score order.
    pub class_labels: Vec<i64>,
    /// Optional first probability calibration coefficients.
    pub prob_a: Vec<f64>,
    /// Optional second probability calibration coefficients.
    pub prob_b: Vec<f64>,
    /// Kernel configuration.
    pub kernel: SvmKernel,
}

impl SvmClassifier {
    /// Validates this raw ONNX-ML classifier representation.
    ///
    /// # Errors
    ///
    /// Returns an error for inconsistent class, support-vector, coefficient,
    /// bias, or probability-calibration lengths.
    pub fn validate(&self) -> Result<()> {
        let class_count = self.class_labels.len();
        let support_count = self.support_vectors.nrows();
        let pair_count = class_count.saturating_mul(class_count.saturating_sub(1)) / 2;
        let valid = class_count >= 2
            && self.support_vectors.ncols() > 0
            && self.vectors_per_class.len() == class_count
            && self.vectors_per_class.iter().sum::<usize>() == support_count
            && self.coefficients.len() == support_count * (class_count - 1)
            && self.rho.len() == pair_count
            && self.prob_a.len() == self.prob_b.len()
            && (self.prob_a.is_empty() || self.prob_a.len() == pair_count);
        if !valid {
            return Err(Error::InvalidModel("SVM classifier shape mismatch".into()));
        }
        Ok(())
    }
}