onnx_pb/
util.rs

1//! Model saving and opening.
2
3use std::path::Path;
4
5use prost::Message;
6
7use crate::ModelProto;
8
9/// File utils error.
10#[derive(Debug)]
11pub enum Error {
12    /// IO error.
13    Io(std::io::Error),
14
15    /// Decode error.
16    Decode(prost::DecodeError),
17
18    /// Encode error.
19    Encode(prost::EncodeError),
20}
21
22/// Opens model from a file.
23pub fn open_model<P: AsRef<Path>>(path: P) -> Result<ModelProto, Error> {
24    let body = std::fs::read(path).map_err(|e| Error::Io(e))?;
25    ModelProto::decode(body.as_slice()).map_err(|e| Error::Decode(e))
26}
27
28/// Saves model to a file.
29pub fn save_model<P: AsRef<Path>>(path: P, model: &ModelProto) -> Result<(), Error> {
30    let mut body = Vec::new();
31    model.encode(&mut body).map_err(|e| Error::Encode(e))?;
32    std::fs::write(path, body).map_err(|e| Error::Io(e))?;
33    Ok(())
34}
onnx_pb/util.rs

onnx_pb/
util.rs
