使用 WebGPU 执行提供程序

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

基础

什么是 WebGPU？我应该使用它吗？

WebGPU 是一个用于通用 GPU 计算和图形的新 Web 标准。它被设计为一个低级 API，基于 D3D12、Vulkan 和 Metal，旨在用于浏览器。它旨在比 WebGL 更高效、性能更好，并用于机器学习、图形和其他计算任务。

WebGPU 在 Windows、macOS、Android 和 ChromeOS 上最新版本的 Chrome 和 Edge 中开箱即用。它也在 Firefox 中通过一个标志启用，并在 Safari Technology Preview 中可用。请查看 WebGPU 状态获取最新信息。

如果您在 Web 应用程序中使用 ONNX Runtime Web 推理非常轻量级的模型，并且希望减小二进制文件大小，您可以继续使用默认的 WebAssembly (WASM) 执行提供程序。如果您想运行计算密集型模型，或者想利用客户端设备中的 GPU，您可以使用 WebGPU 执行提供程序。

如何在 ONNX Runtime Web 中使用 WebGPU EP

本节假设您已经使用 ONNX Runtime Web 设置了您的 Web 应用程序。如果您还没有设置，可以参考开始使用以获取一些基本信息。

要使用 WebGPU EP，您只需进行 2 个小改动：

更新您的导入语句
- 对于 HTML 脚本标签，将 ort.min.js 更改为 ort.webgpu.min.js
```
<script src="https://example.com/path/ort.webgpu.min.js"></script>
```
- 对于 JavaScript 导入语句，将 onnxruntime-web 更改为 onnxruntime-web/webgpu
```
import * as ort from 'onnxruntime-web/webgpu';
```
详见条件导入。

在会话选项中显式指定“webgpu”EP

const session = await ort.InferenceSession.create(modelPath, { ..., executionProviders: ['webgpu'] });

您也可以考虑安装最新每夜构建版本的 ONNX Runtime Web (onnxruntime-web@dev) 以受益于最新的功能和改进。

WebGPU EP 功能

ONNX Runtime Web 提供了以下功能，这些功能可能有助于与 WebGPU EP 配合使用

图捕获

如果您的模型具有静态形状且所有计算内核都在 WebGPU EP 上运行，您可以尝试图捕获功能。此功能可能会潜在地提高模型的性能。

详见图捕获。

使用 `ort.env.webgpu` 标志

详见 env.webgpu。

将张量数据保留在 GPU 上（I/O 绑定）

默认情况下，模型的输入和输出是张量，数据存储在 CPU 内存中。当您使用 WebGPU EP 运行会话时，数据会被复制到 GPU 内存，结果会复制回 CPU 内存。如果您从基于 GPU 的源获取输入数据，或者您想将输出数据保留在 GPU 上以供进一步处理，您可以使用 I/O 绑定将数据保留在 GPU 上。这在运行基于 Transformer 的模型时特别有用，因为这类模型通常会多次运行单个模型，并将前一个输出作为下一个输入。

对于模型输入，如果您的输入数据是 WebGPU 存储缓冲区，您可以创建 GPU 张量并将其用作输入张量。

对于模型输出，有两种方法可以使用 I/O 绑定功能：

使用预分配的 GPU 张量
指定输出数据位置

另请查看以下主题：

零大小张量
GPU 张量生命周期管理

从 GPU 缓冲区创建输入张量

如果您的输入数据是 WebGPU 存储缓冲区，您可以创建一个 GPU 张量并将其用作输入张量

const inputTensor = ort.Tensor.fromGpuBuffer(inputGpuBuffer, {
  dataType: 'float32',
  dims: [1, 3, 224, 224]
});

将此张量用作模型输入 (feeds)，这样输入数据将保留在 GPU 上。

使用预分配的 GPU 张量

如果您事先知道输出形状，您可以创建一个 GPU 张量并将其用作输出张量

// Create a pre-allocated buffer and the corresponding tensor. Assuming that the output shape is [10, 1000].
const bufferSize = (10 * 1000) /* number of elements */ * 4 /* bytes per element */;
const device = ort.env.webgpu.device;
const myPreAllocatedBuffer = device.createBuffer({
    usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST | GPUBufferUsage.STORAGE,
    size: Math.ceil(bufferSize / 16) * 16 /* align to 16 bytes */
});

const myPreAllocatedOutputTensor = ort.Tensor.fromGpuBuffer(myPreAllocatedBuffer, {
  dataType: 'float32',
  dims: [10, 1000]
});

// ...

// Run the session with fetches
const feeds = { 'input_0': myInputTensor };
const fetches = { 'output_0': myPreAllocatedOutputTensor };
const results = await mySession.run(feeds, fetches);

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

指定输出数据位置

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

const mySessionOptions1 = {
  ...,
  // keep all output data on GPU
  preferredOutputLocation: 'gpu-buffer'
};

const mySessionOptions2 = {
  ...,
  // 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': 'gpu-buffer'   // keep output_1 on GPU buffer
  }
};

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

详见 API 参考: preferredOutputLocation。

注意事项

零大小张量

如果张量的形状包含 1 个或多个大小为 0 的维度，则该张量被视为零大小张量。零大小张量不包含任何数据，因此不适用数据位置。ONNX Runtime Web 始终将零大小张量视为 CPU 张量。要创建零大小张量，您可以使用以下代码：

const zeroSizedTensor = new ort.Tensor('float32', [], [3, 256, 0, 64]);

GPU 张量生命周期管理

了解底层 GPU 缓冲区如何管理至关重要，这样可以避免内存泄漏并提高缓冲区使用效率。

GPU 张量由用户代码创建或由 ONNX Runtime Web 作为模型输出创建。

当它由用户代码创建时，它总是使用现有 GPU 缓冲区通过 Tensor.fromGpuBuffer() 创建。在这种情况下，张量不“拥有”GPU 缓冲区。
- 用户有责任确保底层缓冲区在推理期间有效，并在不再需要时调用 buffer.destroy() 来释放缓冲区。
- 避免调用 tensor.getData() 和 tensor.dispose()。直接使用 GPU 缓冲区。
- 使用已销毁 GPU 缓冲区的 GPU 张量将导致会话运行失败。
当它由 ONNX Runtime Web 作为模型输出（不是预分配的 GPU 张量）创建时，张量“拥有”该缓冲区。
- 您无需担心在张量使用前缓冲区被销毁的情况。
- 调用 tensor.getData() 将数据从 GPU 缓冲区下载到 CPU 并以类型化数组形式获取数据。
- 显式调用 tensor.dispose() 以在不再需要时销毁底层 GPU 缓冲区。