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

use crate::canonical::AffineModel;
use crate::graph_builder::{assemble_model, make_node, make_tensor, make_value_info, Dimension};
use crate::proto::{GraphProto, ModelProto};
use crate::{IR_VERSION, OPSET_VERSION};

/// Exports a general affine transform as `Gemm(X, matrix, bias)`.
#[must_use]
pub fn export_affine(model: &AffineModel) -> ModelProto {
    let matrix = model.matrix.mapv(|value| value as f32);
    let bias = model.bias.mapv(|value| value as f32);
    let graph = GraphProto {
        node: vec![make_node(
            "Gemm",
            ["X", "matrix", "bias"],
            ["Y"],
            Vec::new(),
        )],
        name: "affine_transform".into(),
        initializer: vec![
            make_tensor("matrix", &matrix.into_dyn()),
            make_tensor("bias", &bias.into_dyn()),
        ],
        doc_string: String::new(),
        input: vec![make_value_info(
            "X",
            &[
                Dimension::Symbolic("batch".into()),
                Dimension::Fixed(model.n_features()),
            ],
        )],
        output: vec![make_value_info(
            "Y",
            &[
                Dimension::Symbolic("batch".into()),
                Dimension::Fixed(model.n_outputs()),
            ],
        )],
        value_info: Vec::new(),
    };
    assemble_model(graph, OPSET_VERSION, IR_VERSION)
}