使用 WebGPU 执行提供程序

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

目录

• 基础知识
  • 什么是 WebGPU？我应该使用它吗？
  • 如何在 ONNX Runtime Web 中使用 WebGPU EP
• WebGPU EP 功能
  • 图捕获
  • 使用 ort.env.webgpu 标志
• 将张量数据保留在 GPU 上 (IO 绑定)
  • 从 GPU 缓冲区创建输入张量
  • 使用预分配的 GPU 张量
  • 指定输出数据位置
• 注释
  • 零尺寸张量
  • GPU 张量生命周期管理

基础知识

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

WebGPU 是一种用于通用 GPU 计算和图形的新 Web 标准。它被设计为基于 D3D12、Vulkan 和 Metal 的底层 API，旨在在浏览器中使用。它的设计比 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 个小的更改

1. 更新您的导入语句

   • 对于 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';
     

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

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

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

您还可以考虑安装最新 nightly build 版本的 ONNX Runtime Web (onnxruntime-web@dev)，以受益于最新的功能和改进。

WebGPU EP 功能

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

图捕获

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

有关更多详细信息，请参阅图捕获。

使用 ort.env.webgpu 标志

有关更多详细信息，请参阅env.webgpu。

将张量数据保留在 GPU 上 (IO 绑定)

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

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

对于模型输出，有 2 种方法可以使用 IO 绑定功能

• 使用预分配的 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，并以类型化数组的形式获取数据。
  • 在不再需要底层 GPU 缓冲区时，显式调用 tensor.dispose() 以销毁它。