ウェブ AI モデルのテストを強化: WebGPU、WebGL、ヘッドレス Chrome

ジェイソン メイズ
Jason Mayes
Twitter GitHub
フランソワ ボーフォール
François Beaufort
GitHub

うれしいお知らせです。ユーザーのデバイスで ML モデルを直接実行する優れたウェブ AI アプリケーションを作成しました。クラウドに依存することなく、完全にクライアントサイドのウェブブラウザで実行されます。このようなデバイス上の設計により、ユーザーのプライバシーが強化され、パフォーマンスが向上し、コストが大幅に削減されます。

ただし、1 つのハードルがあります。TensorFlow.js モデルは、CPU（WebAssembly）と、より強力な GPU（WebGL と WebGPU を使用）の両方で動作できます。問題は、選択したハードウェアでブラウザテストを一貫して自動化するにはどうすればよいかということです。

整合性を維持することは、実際のユーザーがデバイスで使用するようにデプロイする前に、ML モデルのパフォーマンスを経時的に比較してイテレーションや改善を行ううえで非常に重要です。

GPU を使用して一貫したテスト環境を設定するのは、予想以上に困難な場合があります。このブログ投稿では、Google が直面した問題とその解決方法を共有し、アプリケーションのパフォーマンスを改善できるようにしました。

これはウェブ AI 開発者だけの話ではありません。ウェブゲームやグラフィックに携わっている方にとっても、この投稿は大変貴重です。

自動化ツールボックスの内容

ここでは、以下を使用しています。

  • 環境: NVIDIA T4 または V100 GPU に接続された Linux ベースの Google Colab ノートブック。必要に応じて、Google Cloud（GCP）などの他のクラウド プラットフォームを使用することもできます。
  • ブラウザ: Chrome は WebGPU をサポートしています。WebGPU は WebGL の後継として強力であり、最新の GPU API の進化をウェブにもたらします。
  • 自動化: Puppeteer は、JavaScript を使用してプログラムでブラウザを制御できる Node.js ライブラリです。Puppeteer を使用すれば、Chrome をヘッドレス モードで自動化できます。つまり、ブラウザはサーバー上でインターフェースを表示せずに実行されます。ここでは、従来のフォームではなく、改善された新しいヘッドレス モードを使用します。

環境を確認する

Chrome でハードウェア アクセラレーションがオンになっているかどうかを確認する最善の方法は、アドレスバーに「chrome://gpu」と入力することです。console.log を使用してプログラムで Puppeteer と同等の処理を行うか、完全なレポートを PDF として保存して手動で確認できます。

/* Incomplete example.js */
import puppeteer from 'puppeteer';

// Configure launch parameters: Expands later
const browser = await puppeteer.launch({
  headless: 'new',
  args:  ['--no-sandbox']
});

const page = await browser.newPage();
await page.goto('chrome://gpu');

// Verify: log the WebGPU status or save the GPU report as PDF
const txt = await page.waitForSelector('text/WebGPU');
const status = await txt.evaluate(g => g.parentElement.textContent);
console.log(status);
await page.pdf({ path: './gpu.pdf' });

await browser.close();

chrome://gpu を開きます。結果は次のようになります。

グラフィック機能のステータス
OpenGL: 無効
Vulkan: 無効
WebGL: ソフトウェアのみ。ハードウェア アクセラレーションは利用できません。
WebGL2: ソフトウェアのみ。ハードウェア アクセラレーションは利用できません。
WebGPU: 無効

問題が検出されました。
WebGPU がブロックリストまたはコマンドラインから無効になっています。

スタートとしては良好ではありません。ハードウェアの検出に失敗していることは明らかです。WebGL、WebGL2、WebGPU が基本的に無効になっているか、ソフトウェアのみ。この問題を解決しているのは私たちだけではありません。Chrome の公式サポート チャネル（1）、（2）など、同じような状況にある人々のオンラインでのディスカッションが数多くあります。

WebGPU と WebGL のサポートを有効にする

デフォルトでは、ヘッドレス Chrome は GPU を無効にします。Linux で有効にするには、ヘッドレス Chrome の起動時に次のフラグをすべて適用します。

  • --no-sandbox フラグは Chrome のセキュリティ サンドボックスを無効にします。これにより、ブラウザのプロセスをシステムの他の部分から分離できます。このサンドボックスを使用せずに root として Chrome を実行することはできません。
  • --headless=new フラグは、新しく改善されたヘッドレス モードで Chrome を実行しますが、UI は表示されません。
  • --use-angle=vulkan フラグは、Chrome に ANGLE の Vulkan バックエンドを使用するよう指示します。これにより、OpenGL ES 2/3 呼び出しが Vulkan API 呼び出しに変換されます。
  • --enable-features=Vulkan フラグを指定すると、Chrome での合成とラスタライズ用の Vulkan グラフィック バックエンドが有効になります。
  • --disable-vulkan-surface フラグは、VK_KHR_surface Vulkan インスタンス拡張機能を無効にします。スワップチェーンを使用する代わりに、ビットブリットが画面上の現在のレンダリング結果に使用されます。
  • --enable-unsafe-webgpu フラグは、Linux 版 Chrome で試験運用版の WebGPU API を有効にし、アダプタのブロックリストを無効にします。

これで、これまでに行ったすべての変更を結合します。スクリプトの全文は次のとおりです。

/* Complete example.js */
import puppeteer from 'puppeteer';

// Configure launch parameters
const browser = await puppeteer.launch({
  headless: 'new',
  args: [
    '--no-sandbox',
    '--headless=new',
    '--use-angle=vulkan',
    '--enable-features=Vulkan',
    '--disable-vulkan-surface',
    '--enable-unsafe-webgpu',
  ]
});

const page = await browser.newPage();
await page.goto('chrome://gpu');

// Verify: log the WebGPU status or save the GPU report as PDF
const txt = await page.waitForSelector('text/WebGPU');
const status = await txt.evaluate(g => g.parentElement.textContent);
console.log(status);
await page.pdf({path: './gpu.pdf'});

await browser.close();

スクリプトを再度実行します。WebGPU の問題は検出されず、値は無効からソフトウェアのみに変更されます。

グラフィック機能のステータス
OpenGL: 無効
Vulkan: 無効
WebGL: ソフトウェアのみ。ハードウェア アクセラレーションは利用できません。
WebGL2: ソフトウェアのみ。ハードウェア アクセラレーションは利用できません。
WebGPU: ソフトウェアのみ。ハードウェア アクセラレーションは利用できません。

ただし、ハードウェア アクセラレーションはまだ使用できず、NVIDIA T4 GPU は検出されません。

適切な GPU ドライバをインストールする

Chrome チームの GPU エキスパートと共同で、chrome://gpu の出力をより詳細に調査しました。Linux Colab インスタンスにインストールされるデフォルト ドライバの問題が見つかりました。Vulkan で問題が発生し、次の出力に示すように、Chrome は GL_RENDERER レベルの NVIDIA T4 GPU を検出できません。これにより、ヘッドレス Chrome で問題が発生します。

デフォルトの出力では NVIDIA T4 GPU は検出されません。
ドライバー情報
GL_RENDERER ANGLE（Google、Vulkan 1.3.0（SwiftShader Device（Subzero）（0x0000C0DE））、SwiftShader ドライバ-5.0.0）

互換性のある適切なドライバをインストールすると、問題が解決します。

ドライバのインストール後に出力を更新。
ドライバー情報
GL_RENDERER ANGLE（NVIDIA Corporation、Tesla T4/PCIe/SSE2、OpenGL ES 3.2 NVIDIA 525.105.17）

正しいドライバをインストールするには、セットアップ中に次のコマンドを実行します。最後の 2 行は、NVIDIA ドライバが検出したものの出力を vulkaninfo とともにログに記録するのに役立ちます。

apt-get install -y vulkan-tools libnvidia-gl-525

// Verify the NVIDIA drivers detects along with vulkaninfo
nvidia-smi
vulkaninfo --summary

スクリプトを再度実行すると、次の結果が得られます。🎉

グラフィック機能のステータス
OpenGL: Enabled
Vulkan: Enabled
WebGL: ハードウェアは高速化されましたが、パフォーマンスが低下しています。
WebGL2: ハードウェアは高速化されましたが、パフォーマンスが低下しています。
WebGPU: ハードウェアは高速化されましたが、パフォーマンスが低下しています。

Chrome の実行時に正しいドライバとフラグを使用することで、WebGPU と WebGL がサポートされるようになりました。また、新しいヘッドレス モードもサポートされています。

チームによる調査の舞台裏

リサーチの結果、Google Colab で実行する必要がある環境に適した作業方法は見つかりませんでしたが、他の環境でも機能すると期待する投稿もあり、有望な結果でした。最終的には、2 つの重要な問題があったため、Colab NVIDIA T4 環境での成功を再現できませんでした。

  1. フラグの組み合わせによっては、GPU の検出は可能ですが、実際に GPU を使用することはできません。
  2. サードパーティの実用ソリューションの例では、古い Chrome ヘッドレス バージョンを使用していましたが、ある時点でサポートを終了し、新しいバージョンに移行する予定です。将来に備えるためには、新しいヘッドレス Chrome に対応するソリューションが必要でした。

画像認識用の TensorFlow.js ウェブページの例を実行することで、GPU が十分に活用されていないことを確認しました。これにより、衣服のサンプル（機械学習の「Hello World」のようなもの）を認識するようにモデルをトレーニングしました。

通常のマシンでは、50 トレーニング サイクル（エポック）はそれぞれ 1 秒未満で実行されます。ヘッドレス Chrome をデフォルトの状態で呼び出すと、JavaScript コンソールの出力を Node.js のサーバーサイド コマンドラインに記録して、トレーニング サイクルの実際の速さを確認できます。

予想どおり、各トレーニング エポックに想定よりもかなり長く（数秒）かかりました。これは、Chrome が GPU を使用する代わりに、通常の JS CPU 実行にフォールバックしたことを示しています。

トレーニング エポックは、それよりも遅いペースで進んでいきます。
図 1: 各トレーニング エポックの実行にかかった時間（秒）を示すリアルタイム キャプチャ。

ドライバを修正し、ヘッドレス Chrome 用のフラグの適切な組み合わせを使用してから、TensorFlow.js トレーニング サンプルを再実行すると、トレーニング エポックが大幅に短縮されます。

エポックの速度が向上している..
図 2: エポックの高速化を示すリアルタイム キャプチャ。

概要

ウェブ AI は 2017 年に誕生して以来、飛躍的に成長しました。WebGPU、WebGL、WebAssembly などのブラウザ テクノロジーを使用すると、クライアント側で ML モデルの数学的演算をさらに高速化できます。

2023 年の時点で、TensorFlow.js と MediaPipe Web は 10 億を超えるモデルとライブラリのダウンロードを突破しました。これは歴史的なマイルストーンであり、ウェブ デベロッパーとエンジニアが次世代ウェブアプリに AI を導入し、真に驚くべきソリューションを生み出す傾向を示すものでもあります。

大きな成功には大きな責任が伴います。本番環境システムでのこのレベルの使用では、クライアントサイドのブラウザベースの AI モデルを実際のブラウザ環境でテストするとともに、スケーラブルで自動化可能で、既知の標準化されたハードウェア設定内でテストする必要があります。

新しい Headless Chrome と Puppeteer の力を組み合わせることで、標準化された複製可能な環境でそのようなワークロードを自信を持ってテストし、一貫性のある信頼性の高い結果を得ることができます。

まとめ

Google のドキュメントにある手順ガイドで、完全な設定を自分で試すことができます。

こちらの方法が役に立ったと思ったら、LinkedIn、X（旧 Twitter）、またはハッシュタグ #WebAI を使用して使用しているソーシャル ネットワークで、フィードバックをお寄せください。皆様からのフィードバックをお待ちしております。今後このような投稿を増やしていく予定です。

今後のアップデートを受け取るには、GitHub リポジトリにスターを付ける必要があります。

謝辞

このソリューションで直面したドライバと WebGPU の問題のデバッグに協力してくれた Chrome チームのすべての関係者に、深く感謝します。また、このブログ投稿の文章作成に協力してくれた Jecelyn Yeen と Alexandra White に深く感謝します。最終的な実用的なソリューションの開発に役立った Yuly Novikov、Andrey Kosyakov、Alex Rudenko に感謝します。