1 use anyhow::{Context, Result};
2 use std::fs;
3 use test_programs::nn::{sort_results, witx};
4 
main() -> Result<()>5 pub fn main() -> Result<()> {
6     let model = fs::read("fixture/model.pt")
7         .context("the model file to be mapped to the fixture directory")?;
8     let graph = witx::load(
9         &[&model],
10         witx::GraphEncoding::Pytorch,
11         witx::ExecutionTarget::CPU,
12     )?;
13     let tensor = fs::read("fixture/kitten.tensor")
14         .context("the tensor file to be mapped to the fixture directory")?;
15     let output_buffer = witx::classify(graph, tensor)?;
16     let result = softmax(output_buffer);
17     let top_five = &sort_results(&result)[..5];
18     assert_eq!(top_five[0].class_id(), 281);
19     println!("found results, sorted top 5: {top_five:?}");
20     Ok(())
21 }
22 
softmax(output_tensor: Vec<f32>) -> Vec<f32>23 fn softmax(output_tensor: Vec<f32>) -> Vec<f32> {
24     let max_val = output_tensor
25         .iter()
26         .cloned()
27         .fold(f32::NEG_INFINITY, f32::max);
28 
29     // Compute the exponential of each element subtracted by max_val for numerical stability.
30     let exps: Vec<f32> = output_tensor.iter().map(|&x| (x - max_val).exp()).collect();
31 
32     // Compute the sum of the exponentials.
33     let sum_exps: f32 = exps.iter().sum();
34 
35     // Normalize each element to get the probabilities.
36     exps.iter().map(|&exp| exp / sum_exps).collect()
37 }
38