使用 WebGPU 执行提供程序

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

基础知识

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

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

WebGPU 在最新版本的 Windows、macOS、Android 和 ChromeOS 上的 Chrome 和 Edge 浏览器中开箱即用。它在 Firefox 中可以通过一个标志启用，也在 Safari 技术预览版中可用。请查看 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。

注意事项

零尺寸张量

如果一个张量的形状包含一个或多个尺寸为 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 缓冲区。