rustを使ってonnxを読み込んでみた

https://cipepser.hatenablog.com/entry/protobuf-read-in-rust を参考にrustを使ってonnxを読み込んでみました。

Cargo.toml

protobufを使えるようにする。

[dependencies]
protobuf = { version = "~2.0", features = ["with-bytes"] }

onnxフォーマットのダウンロード

ここからonnxフォーマットを定義したファイルをダウンロードする。 https://github.com/onnx/onnx/blob/master/onnx/onnx.proto

rustで読めるようにする。

$ protoc --rust_out . onnx.proto 

protoc-gen-rust: program not found or is not executable --rust_out: protoc-gen-rust: Plugin failed with status code 1.

この表示が出たら、protobuf-codegenをインストールする

$ cargo install protobuf-codegen

これでonnx.rsが生成される。

rustからファイルを読む。

配布されているsqueezenetのonnxファイルを読み込む

extern crate protobuf;

use std::fs::File;
use std::io::{BufReader};
use protobuf::{CodedInputStream, Message};

//protoc --rust_out . onnx.proto で生成されたonnx.rsを読み込む
mod onnx;
use onnx::ModelProto;

fn main() {
    let file = File::open("../../data/squeezenet/model.onnx").expect("fail to open file");
    let mut buffered_reader = BufReader::new(file);
    let mut cis = CodedInputStream::from_buffered_reader(&mut buffered_reader);

    let mut u = ModelProto::new();
    u.merge_from(&mut cis).expect("fail to merge");

    println!("producer name: {}", u.get_producer_name());
}

実行結果

producer name: onnx-caffe2

なんかうまく行ってそう。