使用 WebNN 执行提供程序

本文档解释了如何在 ONNX Runtime 中使用 WebNN 执行提供程序。

目录

基础知识

什么是 WebNN?我应该使用它吗?

Web 神经网络 (WebNN) API 是一项新的 Web 标准,允许 Web 应用和框架通过设备上的硬件(例如 GPU、CPU 或专用 AI 加速器 (NPU))加速深度神经网络。

WebNN 在 Windows、Linux、macOS、Android 和 ChromeOS 上最新版本的 Chrome 和 Edge 中可用,位于“启用 WebNN API”标志后面。请查看 WebNN 状态以获取最新实现状态。

请参阅 WebNN 算子,了解 WebNN 执行提供程序中算子支持的最新状态。如果 WebNN 执行提供程序支持模型中的大多数算子(不支持的算子会回退到 WASM EP),并且您希望通过利用设备端加速器实现省电、更快的处理和更流畅的性能,请考虑使用 WebNN 执行提供程序。

如何在 ONNX Runtime Web 中使用 WebNN EP

本节假设您已经使用 ONNX Runtime Web 设置了 Web 应用程序。如果您尚未设置,可以按照开始使用获取一些基本信息。

要使用 WebNN EP,您只需进行 3 个小改动

  1. 更新您的导入语句

    • 对于 HTML 脚本标签,将 ort.min.js 更改为 ort.all.min.js
      <script src="https://example.com/path/ort.all.min.js"></script>
      
    • 对于 JavaScript 导入语句,将 onnxruntime-web 更改为 onnxruntime-web/all
      import * as ort from 'onnxruntime-web/all';
      

    有关详细信息,请参阅条件导入

  2. 在会话选项中明确指定“webnn”EP
    const session = await ort.InferenceSession.create(modelPath, { ..., executionProviders: ['webnn'] });
    

    WebNN EP 还提供了一组选项,用于创建不同类型的 WebNN MLContext。

    • deviceType'cpu'|'gpu'|'npu'(默认值为 'cpu'),指定用于 MLContext 的首选设备类型。
    • powerPreference'default'|'low-power'|'high-performance'(默认值为 'default'),指定用于 MLContext 的首选功耗类型。
    • contextMLContext 类型,允许用户将预创建的 MLContext 传递给 WebNN EP,在 I/O 绑定功能中是必需的。如果提供了此选项,则其他选项将被忽略。

    使用 WebNN EP 选项的示例

    const options = {
       executionProviders: [
         {
           name: 'webnn',
           deviceType: 'gpu',
           powerPreference: "default",
         },
       ],
    }
    
  3. 如果是动态形状模型,ONNX Runtime Web 提供了 freeDimensionOverrides 会话选项来覆盖模型的自由维度。有关详细信息,请参阅 freeDimensionOverrides 介绍

WebNN API 和 WebNN EP 正在积极开发中,您可以考虑安装 ONNX Runtime Web 的最新每夜构建版本 (onnxruntime-web@dev),以受益于最新的功能和改进。

在 WebNN MLTensor 上保留张量数据 (I/O 绑定)

默认情况下,模型的输入和输出是张量,其数据存储在 CPU 内存中。当您使用 WebNN EP 并在“gpu”或“npu”设备类型上运行会话时,数据会复制到 GPU 或 NPU 内存,结果再复制回 CPU 内存。不同设备之间以及不同会话之间的内存复制会给推理时间带来很大的开销,WebNN 提供了一种新的不透明设备特定存储类型 MLTensor 来解决此问题。如果您的输入数据来自 MLTensor,或者您希望将输出数据保留在 MLTensor 上以供进一步处理,则可以使用 I/O 绑定将数据保留在 MLTensor 上。这在运行基于 Transformer 的模型时特别有用,因为这类模型通常会将单个模型多次运行,并将上一次的输出作为下一次的输入。

对于模型输入,如果您的输入数据是 WebNN 存储 MLTensor,您可以创建一个 MLTensor 张量并将其用作输入张量

对于模型输出,有两种方式使用 I/O 绑定功能

另请查看以下主题

注意:MLTensor 需要共享的 MLContext 才能进行 I/O 绑定。这意味着 MLContext 应作为 WebNN EP 选项预先创建,并在所有会话中重复利用。

从 MLTensor 创建输入张量

如果您的输入数据是 WebNN 存储 MLTensor,您可以创建一个 MLTensor 张量并将其用作输入张量

// Create WebNN MLContext
const mlContext = await navigator.ml.createContext({deviceType, ...});
// Create a WebNN MLTensor
const inputMLTensor = await mlContext.createTensor({
  dataType: 'float32',
  shape: [1, 3, 224, 224],
  writable: true,
});
// Write data to the MLTensor
const inputArrayBuffer = new Float32Array(1 * 3 * 224 * 224).fill(1.0);
mlContext.writeTensor(inputMLTensor, inputArrayBuffer);

// Create an ORT tensor from the MLTensor
const inputTensor = ort.Tensor.fromMLTensor(inputMLTensor, {
  dataType: 'float32',
  dims: [1, 3, 224, 224],
});

将此张量用作模型输入 (feeds),以便输入数据将保留在 MLTensor 上。

使用预分配的 MLTensor 张量

如果您提前知道输出形状,可以创建一个 MLTensor 张量并将其用作输出张量


// Create a pre-allocated MLTensor and the corresponding ORT tensor. Assuming that the output shape is [10, 1000].
const mlContext = await navigator.ml.createContext({deviceType, ...});
const preallocatedMLTensor = await mlContext.createTensor({
  dataType: 'float32',
  shape: [10, 1000],
  readable: true,
});

const preallocatedOutputTensor = ort.Tensor.fromMLTensor(preallocatedMLTensor, {
  dataType: 'float32',
  dims: [10, 1000],
});

// ...

// Run the session with fetches
const feeds = { 'input_0': inputTensor };
const fetches = { 'output_0': preallocatedOutputTensor };
await session.run(feeds, fetches);

// Read output_0 data from preallocatedMLTensor if need
const output_0 = await mlContext.readTensor(preallocatedMLTensor);
console.log('output_0 value:', new Float32Array(output_0));

通过在 fetches 中指定输出张量,ONNX Runtime Web 将使用预分配的 MLTensor 作为输出张量。如果存在形状不匹配,则 run() 调用将失败。

指定输出数据位置

如果您不想为输出使用预分配的 MLTensor 张量,也可以在会话选项中指定输出数据位置

const sessionOptions1 = {
  ...,
  // keep all output data on MLTensor
  preferredOutputLocation: 'ml-tensor'
};

const sessionOptions2 = {
  ...,
  // alternatively, you can specify the output location for each output tensor
  preferredOutputLocation: {
    'output_0': 'cpu',         // keep output_0 on CPU. This is the default behavior.
    'output_1': 'ml-tensor'   // keep output_1 on MLTensor tensor
  }
};

// ...

// Run the session
const feeds = { 'input_0': inputTensor };
const results = await session.run(feeds);

// Read output_1 data
const output_1 = await results['output_1'].getData();
console.log('output_1 value:', new Float32Array(output_1));

通过指定配置 preferredOutputLocation,ONNX Runtime Web 将把输出数据保留在指定的设备上。

有关详细信息,请参阅API 参考:preferredOutputLocation

注意事项

MLTensor 张量生命周期管理

了解底层 MLTensor 的管理方式非常重要,这样可以避免内存泄漏并提高张量使用效率。

MLTensor 张量可以通过用户代码创建,也可以由 ONNX Runtime Web 作为模型输出创建。

  • 当它由用户代码创建时,总是使用现有的 MLTensor 通过 Tensor.fromMLTensor() 创建。在这种情况下,张量不“拥有”MLTensor。

    • 用户有责任确保底层 MLTensor 在推理过程中有效,并在不再需要时调用 mlTensor.destroy() 来释放 MLTensor。
    • 避免调用 tensor.getData()tensor.dispose()。直接使用 MLTensor 张量。
    • 使用已销毁的 MLTensor 的 MLTensor 张量将导致会话运行失败。
  • 当它由 ONNX Runtime Web 作为模型输出(而不是预分配的 MLTensor 张量)创建时,该张量“拥有”MLTensor。

    • 您无需担心 MLTensor 在张量使用之前被销毁的情况。
    • 调用 tensor.getData() 将数据从 MLTensor 下载到 CPU 并获取为类型化数组。
    • 当不再需要底层 MLTensor 时,请明确调用 tensor.dispose() 来销毁它。