構建基於 WasmEdge 與 WASI-NN 介面的 OpenVINO 的道路分割推理任務 | WasmEdge 0.10.1 系列解讀文章

語言: CN / TW / HK

作者:apepkuss

7月28日,WasmEdge 0.10.1 正式釋出。今天就帶大家詳細瞭解 0.10.1 版本中的 wasi-nn 提案。本篇是 wasi-nn 系列文章的第一篇,下一篇文章將介紹 WasmEdge 對 wasi-nn 提案的優化點。

如今,AI 推理這個詞已經不是什麼陌⽣的詞彙了。從技術層面上講,AI 推理的路徑是輸入資料,呼叫模型,返回結果。這看起來是完美的 serverless 函式,因為 AI 推理有簡單的輸入與輸出,並且是無狀態的。眾所周知,AI 推理是一項計算密集型工作。使用 WebAssembly 和 Rust 可以實現高效能的 AI 推理函式,同時通過 Wasm 保證函式的安全與跨平臺易用性。

最近,WasmEdge Runtime<sup>[1]</sup> 的 0.10.1 版本已經提供了對 WASI-NN<sup>[2]</sup> 介面提案的支援,後端推理引擎部分目前僅支援 Intel OpenVINO<sup>[3]</sup>

除了 OpenVINO 外,WasmEdge Runtime 還支援 TensorFlow推理引擎,但是這兩種模型採用了兩種不同的支援方案,本文著重介紹 wasi-nn。

不過,根據7月份的 WasmEdge 社群會議上公佈的開發計劃,WasmEdge 後續會逐步支援 TensorRT、PyTorch、ONNX Runtime 等後端推理引擎。那麼,如何使用這套新的介面規範來構建基於 WebAssembly 技術的 AI 推理任務?開發流程是什麼樣子的?複雜程度又如何?

在這篇文章我們嘗試通過一個簡單的道路分割 ADAS例子來回答這些問題。以下所涉及的示例程式碼及相關檔案,可前往 WasmEdge-WASINN-examples 程式碼庫檢視下載。也歡迎你新增更多 wasi-nn example。

(本文內容大綱)

WASI-NN 是什麼?

在示例之前,我們簡單介紹一下 WASI-NN<sup>[4]</sup> 介面提案。

實際上,WASI-NN 提案的名字是由兩個部分構成:WASIWebAssembly System Interface 的縮寫,簡單來說, WASI 定義了一組介面規範,從而允許WebAssembly在更細粒度的許可權控制下,安全地執行在非瀏覽器的環境中;NN 代表 Neural Network,即神經網路。顯而易見,WASI-NNWASI 介面規範的一個組成部分,其主要針對機器學習這個應用場景。

理論上來說,這個介面規範既可以用於模型訓練,亦可以用於模型推理,我們的示例僅針對模型推理這個部分。關於 WASIWASI-NN 介面規範的更多細節,可前往 wasi.dev 進一步瞭解。

目前,WasmEdge Runtime 專案的主分支上已經提供了穩定的 WASI-NN 支援,涵蓋了 WASI-NN Proposal Phase 2 所定義的五個主要介面:

// 載入模型的位元組序列
load: function(builder: graph-builder-array, encoding: graph-encoding, target: execution-target) -> expected<graph, error>

// 建立計算圖執行例項
init-execution-context: function(graph: graph) -> expected<graph-execution-context, error>

// 載入輸入
set-input: function(ctx: graph-execution-context, index: u32, tensor: tensor) -> expected<unit, error>

// 執行推理
compute: function(ctx: graph-execution-context) -> expected<unit, error>

// 提取結果
get-output: function(ctx: graph-execution-context, index: u32) -> expected<tensor, error>

這些介面的主要作用就是為實現 Wasm 模組與本地系統資源之間的“互通”提供管道。在推理任務中,前端(Wasm 模組)和後端(推理引擎)之間的資料就是通過這個管道來完成的。下圖是 WasmEdge RuntimeWASI-NN 介面應用簡圖。圖中,綠色矩形所表示的 WASI-NN 介面將前端的 Wasm 模組與後端的 OpenVINO 推理引擎進行了“繫結”。下文示例在執行推理階段,實際上就是通過 WasmEdge Runtime 內建的 WASI-NN 介面,在前、後端之間完成資料互動、函式排程等一系列工作。

下面,就結合具體的例子,來實戰一下如何基於 WasmEdge Runtime 來構建一個“簡約而不簡單”的機器學習推理任務。

使用 OpenVINO 進行道路切割 ADAS

在動手之前,我們先來定義一下使用 WASI-NN 介面構建機器學習推理任務的大致流程,以便對各階段的任務有個總體把握。

  • 任務1:定義推理任務、獲取推理模型和輸入
  • 任務2:環境準備
  • 任務3:構建wasm推理模組
  • 任務4:執行wasm模組。通過 WasmEdge Runtime 提供的命令列執行模式,即 standalone 模式,執行任務3中建立的wasm模組,完成推理任務。
  • 任務5:視覺化推理任務的結果資料。

下面我們就詳細描述一下如何完成上述各項任務目標。

任務1: 推理模型和輸入圖片的獲取

在推理模型的選擇上,為了方便起見,我們在 Intel 官方的 openvino-model-zoo<sup>[5]</sup> 開原始碼庫中選擇了 road-segmentation-adas-0001 模型。這個模型主要用於自動駕駛場景下,完成對道路進行實時分割的任務。為了簡化示例的規模,我們僅使用圖片作為推理任務的輸入。

任務2: 環境準備

我們選擇 Ubuntu 20.04 作為系統環境。WasmEdge 專案也提供了自己的 Ubuntu 20.04 開發環境,所以想簡化環境準備過程的同學,可以從 doker hub 上拉取系統映象。除了系統環境外,還需要部署一下安裝包:

這裡需要說明一下,安裝 Jupyter Notebook 主要出於兩方面的原因:一方面是為了使用Python、Numpy 和 OpenCV 視覺化資料,比如示例圖片和推理結果;另外一方面,通過 Evcxr 外掛,可以獲得一個互動式的輕量級Rust開發環境,很適合用於示例程式碼開發。

環境準備完畢後,就可以下載示例專案的程式碼和相關檔案。本次示例專案的完整程式碼和演示用的相關檔案存放在 WasmEdge-WASINN-examples/openvino-road-segmentation-adas,可以使用下面的命令下載:

// 下載示例專案
git clone [email protected]:second-state/WasmEdge-WASINN-examples.git

// 進入到本次示例的根目錄
cd WasmEdge-WASINN-examples/openvino-road-segmentation-adas/rust

// 檢視示例專案的目錄結構
tree .

示例專案的目錄結構應該是下面這個樣子:

.
├── README.md
├── image
│   └── empty_road_mapillary.jpg ---------------- (示例中用作推理任務輸入的圖片)
├── image-preprocessor   ------------------------ (Rust專案,用於將輸入圖片轉換為OpenVINO tensor)
│   ├── Cargo.lock
│   ├── Cargo.toml
│   └── src
│       └── main.rs
├── model --------------------------------------- (示例中所使用的OpenVINO模型檔案:xml檔案用於描述模型架構,bin檔案存放模型的權重資料)
│   ├── road-segmentation-adas-0001.bin
│   └── road-segmentation-adas-0001.xml
├── openvino-road-segmentation-adas-0001 -------- (Rust專案,其中定義了wasi-nn介面呼叫邏輯。編譯為wasm模組,通過WasmEdge CLI呼叫執行)
│   ├── Cargo.lock
│   ├── Cargo.toml
│   └── src
│       └── main.rs
├── tensor -------------------------------------- ()
│   ├── wasinn-openvino-inference-input-512x896x3xf32-bgr.tensor ---(該二進位制檔案由輸入圖片轉化而來,作為wasm推理模組的一個輸入)
│   └── wasinn-openvino-inference-output-1x4x512x896xf32.tensor  ---(該二進位制檔案儲存了wasm推理模組產生的結果資料
└── visualize_inference_result.ipynb ------------ (用於視覺化資料)

根據上面目錄結構中的註釋,各位同學應該對這個示例專案的各個部分有了大概的瞭解。這裡再說明幾點:

  • 示例圖片轉換為OpenVINO Tensor

    對於初次使用 OpenVINO 作為 WASI-NN介面後端的開發者來說,此處算是一個坑。

    • 第一,Intel 官方在其 openvino-rs<sup>[6]</sup> 開源專案的示例中,使用了 *.bgr 檔案作為推理任務的輸入。這個檔案實際上是一個二進位制檔案,而 bgr 表示檔案資料對應的是 BGR 格式的圖片。另外,在 openvino-rs 專案中可以找到一個名為 openvino-tensor-converter 的工具。這個工具就是用來生成示例中的 *.bgr 檔案。我們示例專案中的 image-preprocessor 也是基於這個工具改進而來的。
    • 第二個容易出錯的地方是輸入 tensor 的維度排布。在Intel官方的 openvino-model-zooopenvino-notebooks<sup>[7]</sup> 開源專案的文件中,均使用了 NCHW 作為輸入tensor的維度排布;並且在使用Python API 和 Rust API 驗證時,也遵從這樣的維度排布。但是,使用 wasi-nn crate<sup>[8]</sup> 時,輸入tensor的維度排布則是 HWC。出現這種情況的具體原因暫時還不確定。
  • image-preprocessoropenvino-road-segmentation-adas-0001 這兩個子專案都是 Rust 專案,沒有把它們整合到一個專案裡的原因在於,前者依賴 opencv-rs<sup>[9]</sup>,導致無法編譯為wasm模組。目前一個值得嘗試的解決辦法是將 opencv-rs 替換為 image<sup>[10]</sup>,感興趣的同學可以嘗試替換一下。

任務3: wasm 推理模組

因為示例的側重點是 WASI-NN 介面,所以對 image-preprocessor 這個部分就不進行過多的介紹,感興趣的同學可以詳細看一下程式碼,應該很快就能理解。那麼,現在我們就來看一下 WASI-NN 介面。WebAssembly.org 在其官方Github上釋出的 WebAssembly/wasi-nn<sup>[11]</sup> 程式碼庫中,提供了兩個比較重要的文件,一個是 wasi-nn.wit.md,一個是 wasi-nn.abi.md。前者使用 wit 語法格式 描述了 WASI-NN 介面規範所涉及的介面及相關資料結構,而後者則是針對前者中所涉及的資料型別給出了更為明確的定義。下面是 wasi-nn.wit.md 給出的五個介面函式:

// 第一步:載入本次推理任務所需要的模型檔案和配置
// builder: 需要載入的模型檔案
// encoding: 後端推理引擎的型別,比如openvino, tensorflow等
// target: 所採用的硬體加速器型別,比如cpu, gpu等
load: function(builder: graph-builder-array, encoding: graph-encoding, target: execution-target) -> expected<graph, error>

// 第二步:通過第一步建立的graph,初始化本次推理任務的執行環境。
// graph-execution-context實際上是對後端推理引擎針對本次推理任務所建立的一個session的封裝,主要的作用就是將第一步中建立的graph和第三步中提供
// 的tensor進行繫結,以便在第四步執行推理任務中使用。
init-execution-context: function(graph: graph) -> expected<graph-execution-context, error>

// 第三步:設定本次推理任務的輸入。
set-input: function(ctx: graph-execution-context, index: u32, tensor: tensor) -> expected<unit, error>

// 第四步:執行本次推理任務
compute: function(ctx: graph-execution-context) -> expected<unit, error>

// 第五步:推理任務成功結束後,提取推理結果資料。
get-output: function(ctx: graph-execution-context, index: u32) -> expected<tensor, error>

從上面的註釋部分可以看出,這五個介面函式構成了使用 WASI-NN 介面完成一次推理任務的模板。因為上述提及的兩份 wit 格式檔案只是給出了 WASI-NN 介面的“形式化”定義,因此每種程式語言可以再進一步例項化這些介面。在 Rust 語言社群, Intel 的兩位工程師 Andrew BrownBrian Jones 共同建立了 WASI-NN 的 Rust binding: wasi-nn crate。我們的示例會通過這個 crate 提供的介面來構建推理模組。

接下來,我們看一下本示例中用於構建推理 Wasm 模組的 openvino-road-segmentation-adas-0001 子專案。下面的程式碼片段是這個專案中最主要的部分:推理函式。

// openvino-road-segmentation-adas-0001/src/.main.rs

/// Do inference
fn infer(
    xml_bytes: impl AsRef<[u8]>,
    weights: impl AsRef<[u8]>,
    in_tensor: nn::Tensor,
) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
    // 第一步:載入本次推理任務所需要的模型檔案和配置
    let graph = unsafe {
        wasi_nn::load(
            &[xml_bytes.as_ref(), weights.as_ref()],
            wasi_nn::GRAPH_ENCODING_OPENVINO,
            wasi_nn::EXECUTION_TARGET_CPU,
        )
        .unwrap()
    };

  	// 第二步:通過第一步建立的graph,初始化本次推理任務的執行環境
    let context = unsafe { wasi_nn::init_execution_context(graph).unwrap() };

  	// 第三步:設定本次推理任務的輸入
    unsafe {
        wasi_nn::set_input(context, 0, in_tensor).unwrap();
    }
  
    // 第四步:執行本次推理任務
    unsafe {
        wasi_nn::compute(context).unwrap();
    }

    // 第五步:推理任務成功結束後,提取推理結果資料
    let mut output_buffer = vec![0f32; 1 * 4 * 512 * 896];
    let bytes_written = unsafe {
        wasi_nn::get_output(
            context,
            0,
            &mut output_buffer[..] as *mut [f32] as *mut u8,
            (output_buffer.len() * 4).try_into().unwrap(),
        )
        .unwrap()
    };

    println!("bytes_written: {:?}", bytes_written);

    Ok(output_buffer)
}

infer 函式體的程式碼邏輯可以發現:

  • 介面函式的呼叫邏輯完全復刻了之前描述的 WASI-NN 介面呼叫模板。
  • 目前 wasi-nn crate 提供的依然是 unsafe 介面。對於 WasmEdge Runtime 社群,在現有 wasi-nn crate的基礎上提供一個安全封裝的crate,對於社群開發者來說,使用起來會更為友好。

因為我們的示例是準備通過 WasmEdge Runtime 提供的命令列介面來執行,所以我們就將 infer 函式所在的 openvino-road-segmentation-adas-0001 子專案編譯為wasm模組。開始編譯前,請通過下面的命令確定 rustup 工具鏈是否安裝了 wasm32-wasi target。

rustup target list

如果在返回結果中沒有看到 wasm32-wasi (installed) 字樣,則可以通過下面的命令安裝:

rustup target add wasm32-wasi

現在可以執行下面的命令,編譯獲得推理 Wasm 模組:

// 確保當前目錄為 openvino-road-segmentation-adas-0001 子專案的根目錄 

cargo build --target=wasm32-wasi --release

如果編譯成功,在 ./target/wasm32-wasi/release 路徑下,可以找到名為 rust-road-segmentation-adas.wasm 的模組,即負責呼叫 WASI-NN 介面的 Wasm 模組。

任務4:執行 wasm 模組

根據 rust-road-segmentation-adas.wasm 模組的入口函式,通過 WasmEdge Runtime 命令列介面呼叫該模組時,需要提供三個輸入(見下面程式碼段中的註釋):

// openvino-road-segmentation-adas-0001/src/main.rs

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args: Vec<String> = env::args().collect();
  	// openvino模型架構檔案
    let model_xml_name: &str = &args[1];
  	// openvino模型權重檔案
    let model_bin_name: &str = &args[2];
  	// 由圖片轉換得到的openvino tensor檔案
    let tensor_name: &str = &args[3];

    ...
}

為了方便復現,可以在示例專案中找到示例所需的檔案:

  • Road-segmentation-adas-0001模型架構檔案:model/road-segmentation-adas-0001.xml
  • Road-segmentation-adas-0001模型權重檔案:model/road-segmentation-adas-0001.bin
  • 推理所用的輸入:tensor/wasinn-openvino-inference-input-512x896x3xf32-bgr.tensor

由於我們借用了 Intel 官方 openvino model zoo 中的模型,所以有關模型的輸入、輸出等資訊可以在road-segmentation-adas-0001模型頁面找到。此外,圖片檔案並不能直接作為輸入,而是需要經過一些預處理,比如 resize 和 RGB 轉 BGR 等,之後再轉換為位元組序列,如此才能通過 wasi-nn crate 提供的介面傳遞給後端的推理引擎。上面的 *.tensor 檔案就是圖片檔案 image/empty_road_mapillary.jpg 經過 image-preprocessor 工具預處理後匯出的二進位制檔案。如果你想推理過程中嘗試使用自己的影象,那麼可以通過下面提供的兩種方式獲得相應的 *.tensor 檔案:

// 進入 image-preprocessor 子專案的根目錄,執行以下命令

cargo run -- --image ../image/empty_road_mapillary.jpg --dims 512x896x3xfp32 --tensor wasinn-openvino-inference-input-512x896x3xf32-bgr.tensor

// 或者,通過編譯image-preprocessor子專案得到im2tensor可執行檔案,再執行轉換

cargo build --release
cd ./target/release
im2tensor --image ../image/empty_road_mapillary.jpg --dims 512x896x3xfp32 --tensor  wasinn-openvino-inference-input-512x896x3xf32-bgr.tensor

在輸入檔案準備好後,我們就可以通過 WasmEdge Runtime 提供的命令列工具來執行推理任務。

  • 首先,確認 WasmEdge Runtime 的命令列工具已經部署到本地系統:

    wasmedge --version
    
    // 或者
    
    /your/local/path/to/wasmedge-release/bin/wasmedge --version
    

    如果你沒有看到 wasmedge version 0.10.0.-71-ge920d6e6 或者類似的版本資訊,那麼你可以按照 WasmEdge Runtime 官方安裝指南 上的步驟完成安裝。

  • 如果 WasmEdge Runtime 命令列工具能夠正確工作,那麼就可以執行下面的命令執行推理任務:

    //在本示例專案的根目錄下執行以下命令
    
    wasmedge --dir .:. /path/to/rust-road-segmentation-adas.wasm ./model/road-segmentation-adas-0001.xml ./model/road-segmentation-adas-0001.bin ./tensor/wasinn-openvino-inference-input-512x896x3xf32-bgr.tensor
    

    推理任務開始執行後,在終端上應該會列印如下資訊:

    Load graph XML, size in bytes: 401509
    Load graph weights, size in bytes: 737192
    Load input tensor, size in bytes: 5505024
    Loaded graph into wasi-nn with ID: 0
    Created wasi-nn execution context with ID: 0
    Executed graph inference
    bytes_written: 7340032
    dump tensor to "wasinn-openvino-inference-output-1x4x512x896xf32.tensor" --- 推理任務完成後,結果資料儲存在該二進位制檔案中
            The size of bytes: 7340032   --------------------------------------- 結果資料的位元組數
    

    說明一下,這裡為了增加輸出檔案的可讀性,我們按照一定的規則硬編碼了匯出檔案的名字,其中 1x4x512x896xf32 用於標識輸出資料的原始維度排布為 NCHW 、資料型別為 float32。這樣做的目的是,在後期對結果資料進行後處理或者視覺化等操作時,便於資料轉換。下面,我們就來實際操作一下,使用 Python、Numpy、OpenCV 這樣的組合,在Jupyter Notebook 上對輸入圖片、推理結果資料、最終結果資料進行視覺化。

任務5:推理任務的資料視覺化

為了便於以更直觀的方式觀察推理過程前後的資料,我們使用 Jupyter Notebook 來搭建一個簡單的資料視覺化工具。下面的三幅圖片是對三個部分資料的視覺化結果:中間的 Segmentation 圖片是來自於推理 Wasm 模組,左右兩幅分別是原始圖片、最終結果圖片。關於資料視覺化相關的程式碼定義在 visualize_inference_result.ipynb ,感興趣的同學可以作為參考改寫成自己需要的樣子,這部分就不進行過多的介紹了。從資料視覺化方面來看,Python 生態圈提供的功能性、便利性要遠好於 Rust 生態圈。

歡迎前往 WasmEdge-WASINN-examples 程式碼庫檢視更多例子,也歡迎你新增更多 wasi-nn example。

總結

本文通過一個簡單的例子,展示瞭如何使用 WasmEdge Runtime 提供的 WASI-NN 介面,構建一個道路分割的機器學習示例。

從這個示例中,我們可以觀察到,與傳統機器學習的方法相比,基於 WebAssembly 技術構建機器學習應用所增加的程式碼規模非常有限、增加的額外程式碼維護成本也很低。但是,在應用方面,這些小幅增加的“成本”卻可以幫助獲得更佳的服務效能。比如,在雲服務的環境下,WebAssembly 可以提供比 docker 快100倍的冷啟動速度,執行的持續時間少 10% ~ 50%,極低的儲存空間。

WASI-NN 提案提供了統一的、標準化的介面規範,使得 WebAssembly 執行時能夠通過單一介面與多種型別的機器學習推理引擎後端進行整合,大大降低了系統整合複雜度和後期維護、升級的成本;同時,這一介面規範也提供了一種抽象,將前、後端的細節對彼此進行了隔離,從而有利於快速構建機器學習應用。隨著 WASI-NN 介面規範的不斷完善以及周邊生態的逐步建立,相信 WebAssembly 技術將會以一種質的方式,改變當前機器學習解決方案的部署和應用方式。

參考文獻

<div id="refer-anchor-1"></div>

[1] WasmEdge Runtime GitHub Repo: https://github.com/WasmEdge/WasmEdge

<div id="refer-anchor-2"></div>

[2] WebAssembly System Interface 提案:https://github.com/WebAssembly/WASI

<div id="refer-anchor-3"></div>

[3] Intel OpenVINO 官網 https://docs.openvino.ai/latest/index.html

<div id="refer-anchor-4"></div>

[4] WebAssembly/wasi-nn 提案:https://github.com/WebAssembly/wasi-nn

<div id="refer-anchor-5"></div>

[5] openvino-model-zoo 程式碼庫:https://github.com/openvinotoolkit/open_model_zoo

<div id="refer-anchor-6"></div>

[6] openvino-rs 專案:https://github.com/intel/openvino-rs

<div id="refer-anchor-7"></div>

[7] openvino-notebooks 專案:https://github.com/openvinotoolkit/openvino_notebooks

<div id="refer-anchor-8"></div>

[8] wasi-nn crate: https://crates.io/crates/wasi-nn

<div id="refer-anchor-9"></div>

[9] opencv crate: https://crates.io/crates/opencv

<div id="refer-anchor-10"></div>

[10] image crate: https://crates.io/crates/image

<div id="refer-anchor-11"></div>

[11] WebAssembly WASI-NN 程式碼庫: https://github.com/WebAssembly/wasi-nn