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

Jason Mayes

François Beaufort

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

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

機械学習モデルを反復処理して改善し、実際のユーザーがデバイスで使用できるようにデプロイする前に、機械学習モデルのパフォーマンスを時間の経過とともに比較するには、一貫性を維持することが重要です。

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

これはウェブ AI デベロッパーだけでなく、ウェブ ゲームやグラフィックを扱っている場合も、この投稿は有益です。

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

使用しているツールは次のとおりです。

• 環境: NVIDIA T4 または V100 GPU に接続された Linux ベースの Google Colab ノートブック。必要に応じて、Google Cloud（GCP）などの他のクラウド プラットフォームを使用できます。
• ブラウザ: Chrome は WebGPU をサポートしています。これは、最新の GPU API の進歩をウェブにもたらす、WebGL の後継となる強力な 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 は基本的に無効であるか、ソフトウェアのみです。この問題は私だけが遭遇しているわけではありません。Chrome の公式サポート チャンネル（1）、（2）など、同様の状況にあるユーザーに関するオンラインでの議論が多数あります。

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

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

• --no-sandbox フラグは、Chrome のセキュリティ サンドボックスを無効にします。これにより、ブラウザ プロセスがシステムの他の部分から分離されます。このサンドボックスなしで Chrome を root として実行することはできません。
• --headless=new フラグを使用すると、新しい改良版のヘッドレス モードで Chrome を実行し、UI を表示せずに実行できます。
• --use-angle=vulkan フラグは、ANGLE に Vulkan バックエンドを使用するように Chrome に指示します。これにより、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 の問題は検出されず、値が無効からソフトウェアのみに変更されます。

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

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

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

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

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

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 で実行する必要がある環境で機能する方法は見つかりませんでしたが、他の環境で機能する有望な投稿がいくつかありました。最終的に、Colab NVIDIA T4 環境でこの成功を再現することはできませんでした。2 つの主な問題がありました。

1. フラグの組み合わせによっては、GPU を検出できますが、実際に GPU を使用できません。
2. サードパーティが提供する機能するソリューションの例では、古いバージョンの Chrome ヘッドレスが使用されています。このバージョンは、新しいバージョンに置き換えられて廃止される予定です。将来に備えて、新しい Headless Chrome に対応したソリューションが求められていました。

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

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

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

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

概要

2017 年に誕生して以来、ウェブ AI は急速に成長しています。WebGPU、WebGL、WebAssembly などのブラウザ技術を使用すると、クライアント側で機械学習モデルの数学演算をさらに高速化できます。

2023 年現在、TensorFlow.js と MediaPipe Web のモデルとライブラリのダウンロード数が 10 億回を超えました。これは歴史的なマイルストーンであり、ウェブ デベロッパーとエンジニアが次世代のウェブアプリで AI を活用して驚異的なソリューションを構築する方向にシフトしていることを示しています。

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

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

まとめ

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

この投稿が役に立ったと思われた方は、LinkedIn、X（旧 Twitter）、またはご利用のソーシャル ネットワークでハッシュタグ #WebAI を使用して、ぜひコメントをお寄せください。今後もこのような記事を作成していくために、ご意見をお聞かせいただければ幸いです。

GitHub リポジトリにスターを付けると、今後の最新情報を受け取ることができます。

謝辞

このソリューションで発生したドライバと WebGPU の問題のデバッグに協力してくれた Chrome チームの皆様に心より感謝いたします。特に、このブログ投稿の文言を練り上げてくれた Jecelyn Yeen 氏と Alexandra White 氏に感謝いたします。最終的なソリューションの作成に貢献した Yuly Novikov、Andrey Kosyakov、Alex Rudenko に感謝します。