neural_engine/src/bin/bench_ort.rs

use ort::session::Session;
use ort::session::builder::GraphOptimizationLevel;
use std::collections::HashMap;

fn main() {
    let _ = ort::init().with_name("test").commit();
    let mut session = Session::builder().unwrap()
        .with_optimization_level(GraphOptimizationLevel::Level3).unwrap()
        .with_intra_threads(4).unwrap()
        .commit_from_file("../model.onnx").unwrap();
        
    let ids: Vec<i64> = vec![1, 2, 3, 4, 5];
    let input_value = ort::value::Tensor::from_array((vec![1, 5], ids)).unwrap();
    
    let start = std::time::Instant::now();
    for _ in 0..100 {
        let input_value = ort::value::Tensor::from_array((vec![1, 5], vec![1, 2, 3, 4, 5])).unwrap();
        let _ = session.run(ort::inputs!["input_ids" => input_value]).unwrap();
    }
    let duration = start.elapsed();
    println!("Time for 100 inferences: {:?}", duration);
}