從 WebAssembly 角度改進 WASI-NN | WASI-NN 系列文章2

在上一篇文章中，我們展示瞭如何使用 OpenVINO 構建一個道路分割的機器學習推理任務。在這個過程中，我們觀察到兩個有趣且值得進一步完善的工作：

• 在示例中使用到了 wasi-nn crate，其為 WASI-NN 提案提供了 Rust 介面實現，從而大大降低了使用 Rust 語言構建基於 WebAssembly 技術的機器學習任務的流程複雜度。不過，wasi-nn crate 提供的介面是 unsafe 的，更適合作為底層API 用於構建更高層的庫。因此，我們可以基於 wasi-nn crate 建立一個提供 safe 介面的庫。
• 在對輸入圖片進行預處理的時候，我們使用到了 opencv crate 。但是，因為 opencv crate 無法編譯為 wasm 模組，所以就不得不將圖片預處理模組獨立出來，單獨作為一個專案來實現。

對於上述兩個觀察，我們嘗試做了初步的嘗試：

• 借鑑 Rust 和 WebAssembly 社群開發者的一些嘗試，我們對 wasi-nn crate 中定義的unsafe 介面進行了抽象和安全封裝，構建了 wasmedge-nn crate 原型。本文的後續部分將演示如何使用 wasmedge-nn crate 替換 wasi-nn crate，重新構建上一篇文章中所使用的道路分割 Wasm 推理模組。
• Rust 社群中著名的影象處理庫之一 image crate 提供了我們所需的圖片預處理的基本能力；此外，由於其是 Rust 原生實現，所以基於這個庫來構建我們需要的影象處理庫是可以編譯為 wasm 模組的。

下面，我們繼續使用道路分割示例，具體演示一下我們的改進方案。

wasmedge-nn crate 的安全介面

在上一篇文章中，我們已經使用了 wasi-nn crate 中定義的五個主要的介面，他們分別對應 WASI-NN 提案中的介面。我們對照著看一下改進後的介面。下圖中，藍色框圖中是我們要使用的 wasmedge-nn crate 的 nn 模組中定義的介面，綠色框圖為相對應的 wasi-nn crate 中定義的介面，箭頭顯示了它們之間的對映關係。關於 wasmedge-nn crate 的設計細節，感興趣的同學可以先行閱讀原始碼，後續我們會在另外一篇文章進行討論，所以這裡就不進行過多的闡述了。

基於wasmedge-nn構建wasm推理模組

接下來，我們就通過程式碼來展示如何使用 wasmedge-nn 提供的介面和相關資料結構，重新實現 wasm 推理模組。

下面的示例程式碼是使用 wasmedge-nn crate 提供的安全介面重新構建的 wasm 推理模組。通過程式碼中的註釋，可以很容易地發現：介面的呼叫順序與使用 wasi-nn 介面的呼叫順序保持一致；而最明顯的不同之處在於，因為 wasmedge-nn 中定義的安全介面，所以示例程式碼中不再有 unsafe 字樣出現。正如在上一篇文章中所闡述，示例程式碼中所展示的介面呼叫順序可以看作一個模板：如果更換一個模型來完成一個新的推理任務，下面的程式碼幾乎不需要任何改動。感興趣的同學可以嘗試使用其它的模型來試試。下面示例的完整程式碼可以在這裡找到。

use std::env;
use wasmedge_nn::{
    cv::image_to_bytes,
    nn::{ctx::WasiNnCtx, Dtype, ExecutionTarget, GraphEncoding, Tensor},
};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args: Vec<String> = env::args().collect();
    let model_xml_name: &str = &args[1];
    let model_bin_name: &str = &args[2];
    let image_name: &str = &args[3];

    // 載入圖片，並轉換為位元組序列
    println!("Load image file and convert it into tensor ...");
    let bytes = image_to_bytes(image_name.to_string(), 512, 896, Dtype::F32)?;
  	
  	// 建立 Tensor 例項，包括資料、維度、型別等資訊
    let tensor = Tensor {
        dimensions: &[1, 3, 512, 896],
        r#type: Dtype::F32.into(),
        data: bytes.as_slice(),
    };
  
    // 建立 WASI-NN Context 例項
    let mut ctx = WasiNnCtx::new()?;

  	// 載入模型檔案及其它推理過程需要的配置資訊
    println!("Load model files ...");
    let graph_id = ctx.load(
        model_xml_name,
        model_bin_name,
        GraphEncoding::Openvino,
        ExecutionTarget::CPU,
    )?;

  	// 初始化執行環境
    println!("initialize the execution context ...");
    let exec_context_id = ctx.init_execution_context(graph_id)?;
		
  	// 為執行環境提供輸入
    println!("Set input tensor ...");
    ctx.set_input(exec_context_id, 0, tensor)?;
		
  	// 執行推理計算
    println!("Do inference ...");
    ctx.compute(exec_context_id)?;
		
  	// 獲取推理計算的結果
    println!("Extract result ...");
    let mut out_buffer = vec![0u8; 1 * 4 * 512 * 896 * 4];
    ctx.get_output(exec_context_id, 0, out_buffer.as_mut_slice())?;
		
  	// 匯出計算結果到指定的二進位制檔案
    println!("Dump result ...");
    dump(
        "wasinn-openvino-inference-output-1x4x512x896xf32.tensor",
        out_buffer.as_slice(),
    )?;

    Ok(())
}

這裡需要說明的是，最後匯出的 .tensor 二進位制檔案用於後續視覺化推理結果資料。由於示例程式碼是通過命令列來執行，在某些環境下（比如Docker）無法直接通過 API 呼叫展示推理結果，所以這裡就只是匯出推理結果。對於其他型別的推理任務，比如使用分類模型，在不需要視覺化顯示的情況下，就可以考慮直接列印分類結果，而無需匯出到檔案。作為參考，這裡我們提供一段Python程式碼（引用自WasmEdge-WASINN-examples/openvino-road-segmentation-adas)，通過讀取匯出的 .tensor 檔案，視覺化推理結果資料。

import matplotlib.pyplot as plt
import numpy as np

# 讀取儲存推理結果的二進位制檔案，並將其轉換為原始維度
data = np.fromfile("wasinn-openvino-inference-output-1x4x512x896xf32.tensor", dtype=np.float32)
print(f"data size: {data.size}")
resized_data = np.resize(data, (1,4,512,896))
print(f"resized_data: {resized_data.shape}, dtype: {resized_data.dtype}")

# 準備用於視覺化的資料
segmentation_mask = np.argmax(resized_data, axis=1)
print(f"segmentation_mask shape: {segmentation_mask.shape}, dtype: {segmentation_mask.dtype}")

# 繪製並顯示
plt.imshow(segmentation_mask[0])

基於 image crate 的影象預處理函式

除了提供安全的介面用於執行推理任務，通過 cv 模組，wasmedge-nn crate 提供了基本的影象預處理函式 image_to_bytes。這個函式的實現借鑑了 image2tensor 開源專案的設計，主要用於將輸入圖片轉換為滿足推理任務要求的位元組序列，在後續步驟中進一步構建 Tensor 變數作為推理模組介面函式的輸入。由於當前的後端僅支援 OpenVINO，影象處理的需求還比較簡單，所以這個 cv 模組僅僅包含了這一個影象預處理函式。

use image::{self, io::Reader, DynamicImage};

// 將圖片檔案轉換為特定尺寸，並轉換為指定型別的位元組序列
pub fn image_to_bytes(
    path: impl AsRef<Path>,
    nheight: u32,
    nwidth: u32,
    dtype: Dtype,
) -> CvResult<Vec<u8>> {
  	// 讀取圖片
    let pixels = Reader::open(path.as_ref())?.decode()?;
  	// 轉換為特定的尺寸
    let dyn_img: DynamicImage = pixels.resize_exact(nwidth, nheight, image::imageops::Triangle);
  	// 轉換為BGR格式
    let bgr_img = dyn_img.to_bgr8();
  
  	// 轉換為指定型別的位元組序列
    let raw_u8_arr: &[u8] = &bgr_img.as_raw()[..];
    let u8_arr = match dtype {
        Dtype::F32 => {
            // Create an array to hold the f32 value of those pixels
            let bytes_required = raw_u8_arr.len() * 4;
            let mut u8_arr: Vec<u8> = vec![0; bytes_required];

            for i in 0..raw_u8_arr.len() {
                // Read the number as a f32 and break it into u8 bytes
                let u8_f32: f32 = raw_u8_arr[i] as f32;
                let u8_bytes = u8_f32.to_ne_bytes();

                for j in 0..4 {
                    u8_arr[(i * 4) + j] = u8_bytes[j];
                }
            }

            u8_arr
        }
        Dtype::U8 => raw_u8_arr.to_vec(),
    };

    Ok(u8_arr)
}

有了安全的 wasmedge-nn crate, 與支援將 OpenCV 編譯成 Wasm 的影象處理庫，使用 Rust 與 WebAssembly 進行 AI 推理就變得非常簡單。接下來只需按照第一篇文章的說明執行 OpenVINO 模型就可以了。

總結

wasi-nn crate 為 Rust 開發者提供了基礎性的底層介面，在使用 WasmEdge Runtime 內建的WASI-NN 支援的場景下，大大降低了介面呼叫的複雜性；在此基礎之上，通過提供安全封裝的介面，wasmedge-nn crate 進一步完善了推理任務的使用者介面定義；同時，通過進一步的抽象，將面向推理任務的前端介面與面向推理引擎的後端介面進行了解耦，從而實現前、後端之間的鬆耦合。

此外，通過 cv 模組提供的、基於 image crate 的影象預處理函式，允許影象預處理模組和推理計算模組編譯在同一個 Wasm模組中，從而實現從原始影象到推理任務的輸入張量、再到推理計算、最後到計算結果匯出的流水線化。

關於 wasmedge-nn crate 的細節，我們會在下一篇文章中進行詳細闡述。感興趣的同學也可以前往 wasmedge-nn GitHub repo 進一步瞭解。我們也歡迎對 WasmEdge + AI感興趣的開發者和研究員反饋你們的意見和建議；同時，也歡迎將你們的實踐經驗和故事分享到我們的 WasmEdge-WASINN-examples 開源專案。謝謝！