コンテンツにスキップ

WebGPUの使用

このドキュメントでは、ONNX RuntimeでWebGPU実行プロバイダーを使用する方法について説明します。

  • TOC

WebGPUは、汎用GPU計算とグラフィックスのための新しいWeb標準です。D3D12、Vulkan、Metalに基づく低レベルAPIとして設計され、ブラウザで使用されるように設計されています。WebGLよりも効率的でパフォーマンスが高く、機械学習、グラフィックス、その他の計算タスクに使用されるように設計されています。

WebGPUは、Windows、macOS、Android、ChromeOSの最新バージョンのChromeとEdgeで標準で利用できます。また、Firefoxではフラグの後ろで、Safari Technology Previewでも利用できます。最新情報についてはWebGPU statusを確認してください。

Webアプリケーションで非常に軽量なモデルの推論にONNX Runtime Webを使用し、小さなバイナリサイズを求める場合は、デフォルトのWebAssembly(WASM)実行プロバイダーを使い続けることができます。より計算集約的なモデルを実行したい場合、またはクライアントデバイスのGPUを活用したい場合は、WebGPU実行プロバイダーを使用できます。

ONNX Runtime WebでWebGPU EPを使用する方法

Section titled “ONNX Runtime WebでWebGPU EPを使用する方法”

このセクションでは、ONNX Runtime WebでWebアプリケーションをすでに設定していることを前提としています。まだの場合は、基本情報についてGet Startedに従ってください。

WebGPU EPを使用するには、2つの小さな変更を行うだけです:

  1. インポート文を更新する:

    • HTMLスクリプトタグの場合、ort.min.jsort.webgpu.min.jsに変更:

      <script src="https://example.com/path/ort.webgpu.min.js"></script>
    • JavaScriptインポート文の場合、onnxruntime-webonnxruntime-web/webgpuに変更:

      import * as ort from 'onnxruntime-web/webgpu';

    詳細についてはConditional Importingを参照してください。

  2. セッションオプションで’webgpu’ EPを明示的に指定:

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

最新の機能と改善の恩恵を受けるために、ONNX Runtime Webの最新ナイトリービルドバージョン(onnxruntime-web@dev)のインストールも検討してください。

ONNX Runtime Webは、WebGPU EPと組み合わせて使用すると便利な以下の機能を提供します:

モデルが静的形状を持ち、すべての計算カーネルがWebGPU EPで実行されている場合、グラフキャプチャ機能を試すことができます。この機能により、モデルのパフォーマンスが向上する可能性があります。

詳細についてはGraph Captureを参照してください。

詳細についてはenv.webgpuを参照してください。

GPU上でテンソルデータを保持する(IOバインディング)

Section titled “GPU上でテンソルデータを保持する(IOバインディング)”

デフォルトでは、モデルの入力と出力はCPUメモリにデータを保持するテンソルです。WebGPU EPでセッションを実行すると、データはGPUメモリにコピーされ、結果はCPUメモリにコピーバックされます。GPUベースのソースから入力データを取得する場合、または出力データをさらなる処理のためにGPU上に保持したい場合は、IOバインディングを使用してデータをGPU上に保持できます。これは、通常、前の出力を次の入力として単一のモデルを複数回実行するトランスフォーマーベースのモデルを実行する際に特に役立ちます。

モデル入力の場合、入力データがWebGPUストレージバッファーの場合、GPUバッファーから入力テンソルを作成できます。

モデル出力の場合、IOバインディング機能を使用する方法は2つあります:

以下のトピックも確認してください:

GPUバッファーから入力テンソルを作成

Section titled “GPUバッファーから入力テンソルを作成”

入力データがWebGPUストレージバッファーの場合、GPUテンソルを作成して入力テンソルとして使用できます:

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

このテンソルをモデル入力(feeds)として使用することで、入力データがGPU上に保持されます。

事前割り当てされたGPUテンソルを使用

Section titled “事前割り当てされたGPUテンソルを使用”

出力形状を事前に知っている場合、GPUテンソルを作成して出力テンソルとして使用できます:

// 事前割り当てされたバッファーと対応するテンソルを作成。出力形状が[10, 1000]であると仮定。
const bufferSize = (10 * 1000) /* 要素数 */ * 4 /* 要素あたりのバイト数 */;
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 /* 16バイトにアライン */
});
const myPreAllocatedOutputTensor = ort.Tensor.fromGpuBuffer(myPreAllocatedBuffer, {
dataType: 'float32',
dims: [10, 1000]
});
// ...
// 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 = {
...,
// すべての出力データをGPU上に保持
preferredOutputLocation: 'gpu-buffer'
};
const mySessionOptions2 = {
...,
// または、各出力テンソルの出力場所を指定
preferredOutputLocation: {
'output_0': 'cpu', // output_0をCPU上に保持。これがデフォルトの動作。
'output_1': 'gpu-buffer' // output_1をGPUバッファー上に保持
}
};

preferredOutputLocation設定を指定することで、ONNX Runtime Webは出力データを指定されたデバイス上に保持します。

詳細についてはAPI reference: preferredOutputLocationを参照してください。

テンソルの形状にサイズ0の次元が1つ以上含まれている場合、そのテンソルはゼロサイズテンソルと見なされます。ゼロサイズテンソルにはデータがないため、データの場所は適用されません。ONNX Runtime Webは常にゼロサイズテンソルをCPUテンソルとして扱います。ゼロサイズテンソルを作成するには、以下のコードを使用できます:

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

GPUテンソルライフサイクル管理

Section titled “GPUテンソルライフサイクル管理”

メモリリークを回避し、バッファー使用効率を向上させるために、基盤となるGPUバッファーがどのように管理されるかを理解することが重要です。

GPUテンソルは、ユーザーコードまたはモデルの出力としてONNX Runtime Webによって作成されます。

  • ユーザーコードによって作成される場合、常にTensor.fromGpuBuffer()を使用して既存のGPUバッファーで作成されます。この場合、テンソルはGPUバッファーを「所有」しません。

    • 推論中に基盤となるバッファーが有効であることを確認し、不要になったときにbuffer.destroy()を呼び出してバッファーを破棄するのはユーザーの責任です。
    • tensor.getData()tensor.dispose()の呼び出しは避けてください。GPUバッファーを直接使用してください。
    • 破棄されたGPUバッファーでGPUテンソルを使用すると、セッション実行が失敗します。
  • モデルの出力としてONNX Runtime Webによって作成される場合(事前割り当てされたGPUテンソルではない)、テンソルはバッファーを「所有」します。

    • テンソルが使用される前にバッファーが破棄される場合について心配する必要はありません。
    • tensor.getData()を呼び出してGPUバッファーからCPUにデータをダウンロードし、型付き配列としてデータを取得します。
    • 不要になったときに基盤となるGPUバッファーを破棄するために、明示的にtensor.dispose()を呼び出します。