onnx_std/
model.rs

1//! Owned model type and ergonomic model I/O (ONNX_RS §3 / §4).
2//!
3//! `onnx-std` deliberately does **not** own an IR: the graph, node, value, tensor
4//! and weight types all come from the shared [`onnx_runtime_ir`] crate
5//! (ONNX_RS §4.1 "Shared Crate"), and the protobuf parse/encode stack is reused
6//! from [`onnx_runtime_loader`] (ONNX_RS §3.4 "Built-in: Protobuf Format").
7//!
8//! What this module adds is an *owned, self-contained* [`Model`] — the loader's
9//! [`onnx_runtime_loader::Model`] only borrows a `Graph`, which is awkward to
10//! pass around — plus the thin [`load_model`] / [`save_model`] entry points that
11//! preserve model-level metadata (`ir_version`, producer, `metadata_props`, …)
12//! across a round-trip. The loader's `load_model` drops that metadata, so we
13//! decode it here from the same bytes.
14
15use std::path::Path;
16use std::sync::Arc;
17
18use onnx_runtime_ir::Graph;
19use onnx_runtime_loader::proto::{
20    ModelProto, decode_model,
21    onnx::{ValueInfoProto, type_proto},
22};
23use onnx_runtime_loader::{
24    Model as EncoderModel, ModelMetadata, WeightStore, encode_model_proto,
25    load_model_bytes_with_weights,
26};
27use prost::Message;
28
29use crate::check::{OnnxChecker, ValidationResult};
30use crate::error::{Error, Result};
31use crate::text::{self, PrintOptions};
32
33/// Generated ONNX IR v13 multi-device/sharding protobuf model types.
34///
35/// These are the wire-authoritative Rust types compiled from the vendored ONNX
36/// schema. Re-exporting them here lets callers construct distributed annotations
37/// without depending directly on the runtime loader's protobuf module.
38pub use onnx_runtime_loader::proto::onnx::{
39    DeviceConfigurationProto, IntIntListEntryProto, NodeDeviceConfigurationProto, ShardedDimProto,
40    ShardingSpecProto, SimpleShardedDimProto, simple_sharded_dim_proto,
41};
42
43/// ONNX-ML opaque type payload (`TypeProto.opaque_type`, field 7).
44pub type OpaqueProto = onnx_runtime_loader::proto::onnx::type_proto::Opaque;
45
46/// An owned ONNX model: the shared-IR [`Graph`] plus the model-level metadata
47/// and (optionally) the live weight store backing external initializers.
48///
49/// Unlike [`onnx_runtime_loader::Model`], which borrows a `&Graph`, this type
50/// owns its graph so it can be returned from [`load_model`], inspected, dumped
51/// to text, validated, and written back out.
52pub struct Model {
53    /// The computation graph, in the shared [`onnx_runtime_ir`] IR.
54    pub graph: Graph,
55    /// Model-level metadata not carried by the [`Graph`] itself.
56    pub metadata: ModelMetadata,
57    /// Live weight store backing any `External` initializers, kept alive so the
58    /// model can be re-saved without re-mmapping. `None` for models built in
59    /// memory with only inline weights.
60    weights: Option<Arc<WeightStore>>,
61    /// Exact schema-level representation for models loaded from protobuf or a
62    /// textual protobuf codec. The runtime graph is an execution projection and
63    /// intentionally does not model every ONNX field; retaining the source proto
64    /// makes standard-library serialization lossless for the full bound schema.
65    source_proto: Option<ModelProto>,
66}
67
68impl Model {
69    /// Wrap an in-memory [`Graph`] as a model with default metadata.
70    ///
71    /// The metadata's `opset_import` is still taken from `graph.opset_imports`
72    /// at save time; only the header fields (`ir_version`, producer, …) default.
73    pub fn new(graph: Graph) -> Self {
74        Self {
75            graph,
76            metadata: ModelMetadata::default(),
77            weights: None,
78            source_proto: None,
79        }
80    }
81
82    /// Wrap a [`Graph`] together with explicit model-level metadata.
83    pub fn with_metadata(graph: Graph, metadata: ModelMetadata) -> Self {
84        Self {
85            graph,
86            metadata,
87            weights: None,
88            source_proto: None,
89        }
90    }
91
92    /// Attach the live weight store backing external initializers.
93    pub fn set_weights(&mut self, weights: Arc<WeightStore>) {
94        self.weights = Some(weights);
95    }
96
97    /// The live weight store, if this model carries one.
98    pub fn weights(&self) -> Option<&Arc<WeightStore>> {
99        self.weights.as_ref()
100    }
101
102    /// Construct a model from the complete generated ONNX protobuf.
103    ///
104    /// The protobuf remains the serialization source of truth, while `graph`
105    /// is populated as the runtime-compatible execution projection.
106    pub fn from_proto(proto: ModelProto) -> Result<Self> {
107        let metadata = metadata_from_proto(&proto);
108        let bytes = proto.encode_to_vec();
109        let loaded = load_model_bytes_with_weights(&bytes, ".").or_else(|_| {
110            let projection = execution_projection(&proto);
111            load_model_bytes_with_weights(&projection.encode_to_vec(), ".")
112        })?;
113        Ok(Self {
114            graph: loaded.0,
115            metadata,
116            weights: Some(loaded.1),
117            source_proto: Some(proto),
118        })
119    }
120
121    /// Return the complete generated ONNX protobuf represented by this model.
122    ///
123    /// Models parsed from protobuf JSON/TextFormat return their exact retained
124    /// schema representation. Programmatically-built models are encoded from the
125    /// shared runtime graph.
126    pub fn to_proto(&self) -> Result<ModelProto> {
127        if let Some(proto) = &self.source_proto {
128            return Ok(proto.clone());
129        }
130        let mut encoder = EncoderModel::new(&self.graph).with_metadata(self.metadata.clone());
131        if let Some(weights) = self.weights() {
132            encoder = encoder.with_weights(weights);
133        }
134        Ok(encode_model_proto(&encoder)?)
135    }
136
137    /// Discard the retained source protobuf and make the mutable runtime graph
138    /// authoritative for future serialization.
139    ///
140    /// Full-spec fields that are not represented by the execution IR (such as
141    /// training information and local function declarations) cannot survive
142    /// this transition.
143    pub fn make_graph_authoritative(&mut self) {
144        self.source_proto = None;
145    }
146
147    pub(crate) fn retained_proto(&self) -> Option<&ModelProto> {
148        self.source_proto.as_ref()
149    }
150
151    /// Render this model as a human-readable textual dump (ONNX_RS §5).
152    pub fn to_text(&self) -> String {
153        text::to_text(self)
154    }
155
156    /// Render this model as text with explicit [`PrintOptions`] (ONNX_RS §5.4).
157    pub fn to_text_with(&self, opts: &PrintOptions) -> String {
158        text::to_text_with(self, opts)
159    }
160
161    /// Parse a model previously rendered in the textual format (ONNX_RS §5.4).
162    pub fn from_text(source: &str) -> Result<Self> {
163        text::from_text(source)
164    }
165
166    /// Validate this model with the default [`OnnxChecker`] (ONNX_RS §8).
167    pub fn validate(&self) -> ValidationResult {
168        OnnxChecker::new().check(self)
169    }
170}
171
172/// Load an ONNX model from a `.onnx` protobuf file (ONNX_RS §3.4).
173///
174/// This reuses the runtime's loader pipeline (parse → build IR → resolve
175/// weights → shape inference) and additionally decodes the model-level metadata
176/// (`ir_version`, producer, `doc_string`, `metadata_props`, graph name) so that
177/// a subsequent [`save_model`] reproduces the header faithfully. External
178/// initializer data is resolved relative to the model file's directory and the
179/// backing memory maps are kept alive on the returned [`Model`].
180pub fn load_model(path: impl AsRef<Path>) -> Result<Model> {
181    let path = path.as_ref();
182    let bytes = std::fs::read(path).map_err(|source| Error::Read {
183        path: path.to_path_buf(),
184        source,
185    })?;
186    let proto = decode_model(&bytes)?;
187    let metadata = metadata_from_proto(&proto);
188    let model_dir = path.parent().unwrap_or_else(|| Path::new("."));
189    let (graph, store) = load_model_bytes_with_weights(&bytes, model_dir).or_else(|_| {
190        load_model_bytes_with_weights(&execution_projection(&proto).encode_to_vec(), model_dir)
191    })?;
192    Ok(Model {
193        graph,
194        metadata,
195        weights: Some(store),
196        source_proto: Some(proto),
197    })
198}
199
200/// Save a [`Model`] to a `.onnx` protobuf file (ONNX_RS §3.4).
201///
202/// Models loaded from protobuf or a textual protobuf codec preserve and write
203/// their complete source proto. Programmatically-built models serialize the
204/// shared execution graph and metadata.
205pub fn save_model(model: &Model, path: impl AsRef<Path>) -> Result<()> {
206    let bytes = model.to_proto()?.encode_to_vec();
207    let path = path.as_ref();
208    std::fs::write(path, bytes).map_err(|source| Error::Write {
209        path: path.to_path_buf(),
210        source,
211    })?;
212    Ok(())
213}
214
215/// Build a runtime-loadable projection without changing the retained proto.
216///
217/// Sparse initializers are first-class in the standard but the execution IR
218/// does not yet have sparse weight storage. Treat them as graph inputs in the
219/// projection so references remain structurally valid; serialization always
220/// uses the untouched source proto.
221fn execution_projection(proto: &ModelProto) -> ModelProto {
222    let mut projection = proto.clone();
223    if let Some(graph) = &mut projection.graph {
224        for sparse in &graph.sparse_initializer {
225            let Some(values) = &sparse.values else {
226                continue;
227            };
228            if values.name.is_empty() || graph.input.iter().any(|input| input.name == values.name) {
229                continue;
230            }
231            graph.input.push(ValueInfoProto {
232                name: values.name.clone(),
233                r#type: Some(onnx_runtime_loader::proto::onnx::TypeProto {
234                    value: Some(type_proto::Value::SparseTensorType(
235                        type_proto::SparseTensor {
236                            elem_type: values.data_type,
237                            shape: Some(onnx_runtime_loader::proto::onnx::TensorShapeProto {
238                                dim: sparse
239                                    .dims
240                                    .iter()
241                                    .map(|&dim| {
242                                        onnx_runtime_loader::proto::onnx::tensor_shape_proto::Dimension {
243                                            value: Some(
244                                                onnx_runtime_loader::proto::onnx::tensor_shape_proto::dimension::Value::DimValue(dim),
245                                            ),
246                                            denotation: String::new(),
247                                        }
248                                    })
249                                    .collect(),
250                            }),
251                        },
252                    )),
253                    denotation: String::new(),
254                }),
255                ..Default::default()
256            });
257        }
258    }
259    projection
260}
261
262/// Project the header fields of a decoded [`ModelProto`] into [`ModelMetadata`].
263///
264/// The runtime loader keeps only `opset_imports` on the `Graph`; everything else
265/// is recovered here so the round-trip is faithful.
266fn metadata_from_proto(proto: &ModelProto) -> ModelMetadata {
267    ModelMetadata {
268        ir_version: proto.ir_version,
269        producer_name: proto.producer_name.clone(),
270        producer_version: proto.producer_version.clone(),
271        domain: proto.domain.clone(),
272        model_version: proto.model_version,
273        doc_string: if proto.doc_string.is_empty() {
274            None
275        } else {
276            Some(proto.doc_string.clone())
277        },
278        graph_name: proto
279            .graph
280            .as_ref()
281            .map(|g| g.name.clone())
282            .unwrap_or_default(),
283        metadata_props: proto
284            .metadata_props
285            .iter()
286            .map(|entry| (entry.key.clone(), entry.value.clone()))
287            .collect(),
288    }
289}
290
291#[cfg(test)]
292mod tests {
293    use super::*;
294    use onnx_runtime_ir::{DataType, Node, NodeId, static_shape};
295
296    /// Build a tiny `Z = Add(X, Y)` model and round-trip it through disk.
297    fn add_graph() -> Graph {
298        let mut g = Graph::new();
299        g.opset_imports.insert(String::new(), 21);
300        let x = g.create_named_value("X", DataType::Float32, static_shape([2, 3]));
301        let y = g.create_named_value("Y", DataType::Float32, static_shape([2, 3]));
302        let z = g.create_named_value("Z", DataType::Float32, static_shape([2, 3]));
303        g.add_input(x);
304        g.add_input(y);
305        // `insert_node` overwrites the id, so the placeholder `NodeId(0)` is fine.
306        let mut node = Node::new(NodeId(0), "Add", vec![Some(x), Some(y)], vec![z]);
307        node.name = "add0".to_string();
308        g.insert_node(node);
309        g.add_output(z);
310        g
311    }
312
313    #[test]
314    fn new_wraps_graph_with_default_metadata() {
315        let model = Model::new(add_graph());
316        assert_eq!(model.metadata, ModelMetadata::default());
317        assert_eq!(model.graph.num_nodes(), 1);
318    }
319
320    #[test]
321    fn save_then_load_round_trips_structure_and_metadata() {
322        let meta = ModelMetadata {
323            producer_name: "onnx-std-test".to_string(),
324            graph_name: "g".to_string(),
325            metadata_props: vec![("author".to_string(), "deckard".to_string())],
326            ..Default::default()
327        };
328        let model = Model::with_metadata(add_graph(), meta.clone());
329
330        let dir = std::env::current_dir().unwrap().join("target");
331        std::fs::create_dir_all(&dir).unwrap();
332        let path = dir.join("onnx_std_roundtrip_test.onnx");
333        save_model(&model, &path).unwrap();
334
335        let loaded = load_model(&path).unwrap();
336        assert_eq!(loaded.graph.num_nodes(), 1);
337        assert_eq!(loaded.metadata.producer_name, "onnx-std-test");
338        assert_eq!(loaded.metadata.graph_name, "g");
339        assert_eq!(
340            loaded.metadata.metadata_props,
341            vec![("author".to_string(), "deckard".to_string())]
342        );
343
344        let _ = std::fs::remove_file(&path);
345    }
346
347    #[test]
348    fn load_missing_file_is_read_error() {
349        let result = load_model("definitely-not-a-real-file.onnx");
350        assert!(matches!(result, Err(Error::Read { .. })));
351    }
352}
onnx_std/model.rs

onnx_std/
model.rs
