使用 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 緩衝區。