[Examples] Add the named model feature into pytorch demo. (#27)

Signed-off-by: vincent <vincent@secondstate.io>

pytorch-mobilenet-image/README.md

@@ -10,7 +10,7 @@ This crate depends on the `wasi-nn` in the `Cargo.toml`:
 
 ```toml
 [dependencies]
-wasi-nn = "0.4.0"
+wasi-nn = "0.6.0"
 ```
 
 ## Build
@@ -21,11 +21,13 @@ Compile the application to WebAssembly:
 cargo build --target=wasm32-wasi --release
 ```
 
-The output WASM file will be at [`target/wasm32-wasi/release/wasmedge-wasinn-example-mobilenet-image.wasm`](wasmedge-wasinn-example-mobilenet-image.wasm).
+Because here, we will demonstrate two ways of using wasi-nn. So the output WASM files will be at [`target/wasm32-wasi/release/wasmedge-wasinn-example-mobilenet-image.wasm`](wasmedge-wasinn-example-mobilenet-image.wasm) and [`target/wasm32-wasi/release/wasmedge-wasinn-example-mobilenet-image-named-model.wasm`](wasmedge-wasinn-example-mobilenet-image-named-model.wasm).
 To speed up the image processing, we can enable the AOT mode in WasmEdge with:
 
 ```bash
 wasmedgec rust/target/wasm32-wasi/release/wasmedge-wasinn-example-mobilenet-image.wasm wasmedge-wasinn-example-mobilenet-image-aot.wasm
+
+wasmedgec rust/target/wasm32-wasi/release/wasmedge-wasinn-example-mobilenet-image-named-model.wasm wasmedge-wasinn-example-mobilenet-image-named-model-aot.wasm
 ```
 
 ## Run
@@ -60,6 +62,8 @@ Users should [install the WasmEdge with WASI-NN PyTorch backend plug-in](https:/
 
 Execute the WASM with the `wasmedge` with PyTorch supporting:
 
+- Case 1:
+
 ```bash
 wasmedge --dir .:. wasmedge-wasinn-example-mobilenet-image.wasm mobilenet.pt input.jpg
 ```
@@ -77,3 +81,25 @@ Executed graph inference
    4.) [950](10.4899)orange
    5.) [953](9.4834)pineapple, ananas
 ```
+
+- Case 2: Apply named model feature
+> requirement wasi-nn >= 0.5.0 and WasmEdge-plugin-wasi_nn-(*) >= 0.13.4 and  
+> --nn-preload argument format follow <name>:<encoding>:<target>:<model_path>
+
+```bash
+wasmedge --dir .:. --nn-preload demo:PyTorch:CPU:mobilenet.pt wasmedge-wasinn-example-mobilenet-image-named-model.wasm demo input.jpg
+```
+
+You will get the same output:
+
+```console
+Loaded graph into wasi-nn with ID: 0
+Created wasi-nn execution context with ID: 0
+Read input tensor, size in bytes: 602112
+Executed graph inference
+   1.) [954](20.6681)banana
+   2.) [940](12.1483)spaghetti squash
+   3.) [951](11.5748)lemon
+   4.) [950](10.4899)orange
+   5.) [953](9.4834)pineapple, ananas
+```

pytorch-mobilenet-image/rust/Cargo.toml

@@ -8,6 +8,14 @@ publish = false
 
 [dependencies]
 image = { version = "0.23.14", default-features = false, features = ["gif", "jpeg", "ico", "png", "pnm", "tga", "tiff", "webp", "bmp", "hdr", "dxt", "dds", "farbfeld"]  }
-wasi-nn = { version = "0.4.0" }
+wasi-nn = { version = "0.6.0" }
+
+[[bin]]
+name = "wasmedge-wasinn-example-mobilenet-image"
+path = "src/main.rs"
+
+[[bin]]
+name = "wasmedge-wasinn-example-mobilenet-image-named-model"
+path = "src/named_model.rs"
 
 [workspace]

pytorch-mobilenet-image/rust/src/named_model.rs

@@ -0,0 +1,96 @@
+use image;
+use std::env;
+use std::fs::File;
+use std::io::Read;
+use wasi_nn;
+mod imagenet_classes;
+
+pub fn main() {
+    let args: Vec<String> = env::args().collect();
+    let model_name: &str = &args[1];
+    let image_name: &str = &args[2];
+
+    let graph = wasi_nn::GraphBuilder::new(
+        wasi_nn::GraphEncoding::Autodetec,
+        wasi_nn::ExecutionTarget::CPU,
+    )
+    .build_from_cache(model_name)
+    .unwrap();
+    println!("Loaded graph into wasi-nn with ID: {:?}", graph);
+
+    let mut context = graph.init_execution_context().unwrap();
+    println!("Created wasi-nn execution context with ID: {:?}", context);
+
+    // Load a tensor that precisely matches the graph input tensor (see
+    let tensor_data = image_to_tensor(image_name.to_string(), 224, 224);
+    println!("Read input tensor, size in bytes: {}", tensor_data.len());
+    context
+        .set_input(0, wasi_nn::TensorType::F32, &[1, 3, 224, 224], &tensor_data)
+        .unwrap();
+    // Execute the inference.
+    context.compute().unwrap();
+    println!("Executed graph inference");
+    // Retrieve the output.
+    let mut output_buffer = vec![0f32; 1000];
+    context.get_output(0, &mut output_buffer).unwrap();
+
+    let results = sort_results(&output_buffer);
+    for i in 0..5 {
+        println!(
+            "   {}.) [{}]({:.4}){}",
+            i + 1,
+            results[i].0,
+            results[i].1,
+            imagenet_classes::IMAGENET_CLASSES[results[i].0]
+        );
+    }
+}
+
+// Sort the buffer of probabilities. The graph places the match probability for each class at the
+// index for that class (e.g. the probability of class 42 is placed at buffer[42]). Here we convert
+// to a wrapping InferenceResult and sort the results.
+fn sort_results(buffer: &[f32]) -> Vec<InferenceResult> {
+    let mut results: Vec<InferenceResult> = buffer
+        .iter()
+        .enumerate()
+        .map(|(c, p)| InferenceResult(c, *p))
+        .collect();
+    results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
+    results
+}
+
+// Take the image located at 'path', open it, resize it to height x width, and then converts
+// the pixel precision to FP32. The resulting BGR pixel vector is then returned.
+fn image_to_tensor(path: String, height: u32, width: u32) -> Vec<u8> {
+    let mut file_img = File::open(path).unwrap();
+    let mut img_buf = Vec::new();
+    file_img.read_to_end(&mut img_buf).unwrap();
+    let img = image::load_from_memory(&img_buf).unwrap().to_rgb8();
+    let resized =
+        image::imageops::resize(&img, height, width, ::image::imageops::FilterType::Triangle);
+    let mut flat_img: Vec<f32> = Vec::new();
+    for rgb in resized.pixels() {
+        flat_img.push((rgb[0] as f32 / 255. - 0.485) / 0.229);
+        flat_img.push((rgb[1] as f32 / 255. - 0.456) / 0.224);
+        flat_img.push((rgb[2] as f32 / 255. - 0.406) / 0.225);
+    }
+    let bytes_required = flat_img.len() * 4;
+    let mut u8_f32_arr: Vec<u8> = vec![0; bytes_required];
+
+    for c in 0..3 {
+        for i in 0..(flat_img.len() / 3) {
+            // Read the number as a f32 and break it into u8 bytes
+            let u8_f32: f32 = flat_img[i * 3 + c] as f32;
+            let u8_bytes = u8_f32.to_ne_bytes();
+
+            for j in 0..4 {
+                u8_f32_arr[((flat_img.len() / 3 * c + i) * 4) + j] = u8_bytes[j];
+            }
+        }
+    }
+    return u8_f32_arr;
+}
+
+// A wrapper for class ID and match probabilities.
+#[derive(Debug, PartialEq)]
+struct InferenceResult(usize, f32);