kn_graph/onnx/
mod.rs

1use std::path::Path;
2
3use external_data::ExternalDataLoader;
4pub use load::{InputShaper, OnnxDimValue};
5pub use loader::GraphLoader;
6
7use crate::graph::Graph;
8use crate::onnx::result::OnnxResult;
9use crate::shape::Size;
10
11pub mod external_data;
12mod inputs;
13mod load;
14#[allow(warnings)]
15mod proto;
16pub mod result;
17mod store;
18mod typed_value;
19mod loader;
20
21/// Load an [ONNX](https://github.com/onnx/onnx/blob/main/docs/IR.md) file from the given path.
22///
23/// If `allow_external` is true, the onnx will be allowed to load external data files,
24/// see [the spec](https://github.com/onnx/onnx/blob/main/docs/IR.md#external-tensor-data).
25/// If `allow_external` is false and the ONNX file does reference external data, an error is returned.
26///
27/// For more flexibility, see [GraphLoader].
28pub fn load_graph_from_onnx_path(path: impl AsRef<Path>, allow_external: bool) -> OnnxResult<Graph> {
29    let mut loader = GraphLoader::from_path(path, allow_external)?;
30    loader.add_named_axis("batch_size", Size::BATCH);
31    loader.load()
32}
33
34/// Load an [ONNX](https://github.com/onnx/onnx/blob/main/docs/IR.md) file from the given bytes.
35///
36/// The file is not allowed to reference external data files.
37///
38/// For more flexibility, see [GraphLoader].
39pub fn load_graph_from_onnx_bytes(buffer: &[u8]) -> OnnxResult<Graph> {
40    let mut loader = GraphLoader::from_bytes(buffer);
41    loader.add_named_axis("batch_size", Size::BATCH);
42    loader.load()
43}
44
45pub fn load_graph_from_onnx_bytes_custom(buffer: &[u8], external: Box<dyn ExternalDataLoader>) -> OnnxResult<Graph> {
46    let mut loader = GraphLoader::from_bytes(buffer);
47    loader.set_external_data(external);
48    loader.add_named_axis("batch_size", Size::BATCH);
49    loader.load()
50}
kn_graph/onnx/mod.rs

kn_graph/onnx/
mod.rs
