onnx-export-rs 0.1.1

Export canonical Rust machine-learning models to ONNX
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

//! Tract-backed ONNX round-trip validation.

use std::io::Cursor;

use ndarray::{Array1, Array2};
use tract_onnx::prelude::*;

use crate::{Error, Result};

/// Summary of a prediction comparison.
#[derive(Clone, Debug, PartialEq)]
pub struct ValidationReport {
    /// Largest absolute element-wise difference.
    pub max_absolute_difference: f64,
    /// Mean absolute element-wise difference.
    pub mean_absolute_difference: f64,
    /// Configured tolerance.
    pub tolerance: f64,
    /// Whether all values were finite, shapes matched, and max error passed.
    pub passed: bool,
}

/// Compares source-model output with output produced by an ONNX runtime.
#[must_use]
pub fn compare_predictions(original: &[f64], exported: &[f32], tolerance: f64) -> ValidationReport {
    if original.len() != exported.len() || original.is_empty() || tolerance < 0.0 {
        return ValidationReport {
            max_absolute_difference: f64::INFINITY,
            mean_absolute_difference: f64::INFINITY,
            tolerance,
            passed: false,
        };
    }
    let differences: Vec<_> = original
        .iter()
        .zip(exported)
        .map(|(&left, &right)| (left - f64::from(right)).abs())
        .collect();
    let finite = differences.iter().all(|value| value.is_finite());
    let max = differences.iter().copied().fold(0.0_f64, f64::max);
    let mean = differences.iter().sum::<f64>() / differences.len() as f64;
    ValidationReport {
        max_absolute_difference: max,
        mean_absolute_difference: mean,
        tolerance,
        passed: finite && max <= tolerance,
    }
}

/// Loads and runs exported ONNX bytes with Tract, comparing them to predictions
/// from the source model.
///
/// # Errors
///
/// Returns an error when the model cannot be parsed, optimized, run, or when
/// its first output is not an `f32` tensor.
pub fn validate_export(
    onnx_bytes: &[u8],
    test_inputs: &Array2<f64>,
    original_predict_fn: impl Fn(&Array2<f64>) -> Array1<f64>,
    tolerance: f64,
) -> Result<ValidationReport> {
    let mut reader = Cursor::new(onnx_bytes);
    let runnable = tract_onnx::onnx()
        .model_for_read(&mut reader)
        .and_then(|model| model.into_optimized())
        .and_then(|model| model.into_runnable())
        .map_err(|error| Error::Validation(error.to_string()))?;
    let input_values: Vec<f32> = test_inputs.iter().map(|&value| value as f32).collect();
    let input = Tensor::from_shape(test_inputs.shape(), &input_values)
        .map_err(|error| Error::Validation(error.to_string()))?;
    let outputs = runnable
        .run(tvec!(input.into()))
        .map_err(|error| Error::Validation(error.to_string()))?;
    let exported = outputs[0]
        .to_plain_array_view::<f32>()
        .map_err(|error| Error::Validation(error.to_string()))?;
    let exported: Vec<f32> = exported.iter().copied().collect();
    let original = original_predict_fn(test_inputs);
    Ok(compare_predictions(
        original.as_slice().unwrap_or(&[]),
        &exported,
        tolerance,
    ))
}