Skip to content
This repository has been archived by the owner on Jun 27, 2023. It is now read-only.

yanghaku/wasi-nn-safe

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

14 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

wasi-nn-safe

CI status crates.io status doc.rs status

Introduction

This library provides some convenient and safe wrapper APIs for wasi-nn system calls, which can replace the unsafe wasi-nn APIs.

[dependencies]
wasi-nn-safe = "0.1"

Quick Start

use wasi_nn_safe::{ExecutionTarget, GraphBuilder, GraphEncoding, TensorType};

fn test(model_path: &'static str) -> Result<(), wasi_nn_safe::Error> {
    // prepare input and output buffer.
    let input = vec![0f32; 224 * 224 * 3];
    let input_dim = vec![1, 224, 224, 3];
    // the input and output buffer can be any sized type, such as u8, f32, etc.
    let mut output_buffer = vec![0f32; 1001];

    // build a tflite graph from file.
    let graph = GraphBuilder::new(GraphEncoding::TensorflowLite, ExecutionTarget::CPU)
        .build_from_files([model_path])?;
    // init graph execution context for this graph.
    let mut ctx = graph.init_execution_context()?;
    // set input
    ctx.set_input(0, TensorType::F32, &input_dim, &input)?;
    // do inference
    ctx.compute()?;
    // copy output to buffer
    let output_bytes = ctx.get_output(0, &mut output_buffer)?;

    assert_eq!(output_bytes, output_buffer.len() * std::mem::size_of::<f32>());
    Ok(())
}

Note

This crate is experimental and will change to adapt the upstream wasi-nn specification.

Now version is based on git commit 0f77c48ec195748990ff67928a4b3eef5f16c2de

Related Links

License

This project is licensed under the Apache 2.0 license. See LICENSE for more details.