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

Jason Mayes

François Beaufort

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

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

整合性の維持は、ML モデルのパフォーマンスを経時的に比較して 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 を開くと、次の結果が表示されます。

良いスタートではありません。ハードウェアの検出が失敗していることはよくわかります。 WebGL、WebGL2、WebGPU は基本的に無効になっているか、ソフトウェアのみです。この問題は Google だけではありません。Chrome の公式サポート チャネル（1）、（2）など、オンラインでも同様の状況の人々が議論されています。

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

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

• --no-sandbox フラグは、Chrome のセキュリティ サンドボックスを無効にします。このサンドボックスは、ブラウザ プロセスをシステムの他の部分から分離します。このサンドボックスを使用せずに、root として Chrome を実行することはサポートされていません。
• --headless=new フラグ: UI を表示せずに、改善された新しいヘッドレス モードで Chrome を実行します。
• --use-angle=vulkan フラグは、Chrome に ANGLE の Vulkan バックエンドを使用するように指示します。これにより、OpenGL ES 2/3 呼び出しが Vulkan API 呼び出しに変換されます。
• --enable-features=Vulkan フラグを指定すると、Chrome での合成とラスタライズのために Vulkan グラフィック バックエンドが有効になります。
• --disable-vulkan-surface フラグは、VK_KHR_surface Vulkan インスタンスの拡張機能を無効にします。画面上にレンダリング結果を表示する際、スワップチェーンではなく Bit blit を使用します。
• --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 の問題は検出されず、値は無効からソフトウェアのみに変更されます。

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

正しい GPU ドライバをインストールする

Google は、Chrome チームの GPU エキスパートとともに chrome://gpu の出力をより詳しく調査しました。Linux Colab インスタンスにインストールされているデフォルトのドライバに問題があることが判明しました。これにより、Vulkan で問題が発生し、Chrome が GL_RENDERER レベルで NVIDIA T4 GPU を検出できないという問題が次の出力で検出されました。このため、ヘッドレス Chrome で問題が発生します。

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

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

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

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

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

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

舞台裏: チームの調査

多くの調査の結果、Google Colab で実行する必要がある環境に適した作業方法は見つかっていませんでした。ただし、他の環境でうまく機能する期待に満ちた投稿もいくつかありましたが、これは有望でした。次のような 2 つの重要な問題があったため、最終的には、Colab NVIDIA T4 環境で成功を再現できませんでした。

1. 一部のフラグの組み合わせでは GPU を検出できますが、実際に GPU を使用することはできません。
2. サードパーティによる実用的なソリューションの例では、古い Chrome ヘッドレス バージョンが使用されており、将来的にはサポートが終了し、新しいバージョンへの移行が予定されています。将来の変化への対応を強化するために、新しいヘッドレス Chrome と連携するソリューションが必要でした。

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

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

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

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

まとめ

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

2023 年の時点で、TensorFlow.js と MediaPipe Web はモデルとライブラリのダウンロード数が 10 億回を超えました。これは歴史的なマイルストーンであり、ウェブ デベロッパーやエンジニアが次世代のウェブアプリに AI を導入し、真に素晴らしいソリューションを作り出していることを示しています。

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

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

まとめ

Google のドキュメントに手順ガイドが用意されているので、自分で設定全体を試すことができます。

これが役に立つ場合は、LinkedIn、X（旧 Twitter）、またはハッシュタグ #WebAI を使ってご利用のソーシャル ネットワークにお知らせください。皆様からのご意見、ご感想をお待ちしております。今後の参考にさせていただきます。

今後のアップデートを受け取るには、GitHub リポジトリにスターを追加してください。

謝辞

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