构建基于 WasmEdge 与 WASI-NN 接口的 OpenVINO 的道路分割推理任务

作者：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 提案的名字是由两个部分构成：WASI 是 WebAssembly System Interface 的缩写，简单来说， WASI 定义了一组接口规范，从而允许WebAssembly在更细粒度的权限控制下，安全地运行在非浏览器的环境中；NN 代表 Neural Network，即神经网络。显而易见，WASI-NN 是 WASI 接口规范的一个组成部分，其主要针对机器学习这个应用场景。

理论上来说，这个接口规范既可以用于模型训练，亦可以用于模型推理，我们的示例仅针对模型推理这个部分。关于 WASI 及 WASI-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 Runtime 的 WASI-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 上拉取系统镜像。除了系统环境外，还需要部署一下安装包：

安装 OpenVINO 2021.4 官方安装指南
安装 WasmEdge Runtime v0.10.0 官方安装指南
安装 OpenCV 4.2 安装指南
可选安装 Jupyter Notebook 官方安装指南
安装Rustup编译工具链官方安装指南

这里需要说明一下，安装 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-zoo 和 openvino-notebooks<sup>[7]</sup> 开源项目的文档中，均使用了 NCHW 作为输入tensor的维度排布；并且在使用Python API 和 Rust API 验证时，也遵从这样的维度排布。但是，使用 wasi-nn crate<sup>[8]</sup> 时，输入tensor的维度排布则是 HWC。出现这种情况的具体原因暂时还不确定。
image-preprocessor 和 openvino-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 Brown 和 Brian 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 技术将会以一种质的方式，改变当前机器学习解决方案的部署和应用方式。