ブラウザに AI モデルをキャッシュに保存する

ほとんどの AI モデルには、インターネット経由で転送されるリソースとしてはかなり大きいという共通点があります。最も小さい MediaPipe オブジェクト検出モデル(SSD MobileNetV2 float16)のサイズは 5.6 MB、最大サイズは 25 MB 前後です。

オープンソースの LLM gemma-2b-it-gpu-int4.bin は 1.35 GB で、LLM としては非常に小さいと考えられます。生成 AI モデルは非常に大きくなる可能性があります。そのため、現在、多くの AI がクラウドで使用されています。高度に最適化されたモデルをデバイス上で直接実行するアプリが増えています。ブラウザで実行される LLM のデモはありますが、ブラウザで実行される他のモデルの本番環境グレードの例をいくつか示します。

AI を活用したオブジェクト選択ツールが開いている Adobe Photoshop on the web で、2 頭のジラフと月という 3 つのオブジェクトが選択されています。

今後のアプリの起動を高速化するには、暗黙的な HTTP ブラウザ キャッシュに依存するのではなく、モデルデータをデバイス上に明示的にキャッシュに保存する必要があります。

このガイドでは gemma-2b-it-gpu-int4.bin model を使用して chatbot を作成しますが、このアプローチは、デバイス上の他のモデルやユースケースに合わせて一般化できます。アプリをモデルに接続する最も一般的な方法は、他のアプリリソースとともにモデルを提供することです。配信を最適化することが重要です。

適切なキャッシュ ヘッダーを構成する

サーバーから AI モデルを提供する場合は、正しい Cache-Control ヘッダーを構成することが重要です。次の例は、アプリのニーズに合わせて構築できる堅牢なデフォルト設定を示しています。

Cache-Control: public, max-age=31536000, immutable

リリースされた AI モデルの各バージョンは静的リソースです。変更されないコンテンツには、リクエスト URL でキャッシュ バスティングと組み合わせて長い max-age を指定する必要があります。モデルを更新する必要がある場合は、新しい URL を指定する必要があります。

ユーザーがページを再読み込みすると、サーバーはコンテンツが安定していることを認識しているにもかかわらず、クライアントは再検証リクエストを送信します。immutable ディレクティブは、コンテンツが変更されないため、再検証が不要であることを明示的に示します。immutable ディレクティブは、ブラウザや中間キャッシュ / プロキシ サーバーでは広くサポートされていませんが、汎用性の高い max-age ディレクティブと組み合わせることで、最大限の互換性を確保できます。public レスポンス ディレクティブは、レスポンスを共有キャッシュに保存できることを示します。

Chrome DevTools には、AI モデルをリクエストしたときに Hugging Face から送信された本番環境の Cache-Control ヘッダーが表示されます。(出典

AI モデルをクライアントサイドにキャッシュに保存する

AI モデルを提供する場合は、ブラウザでモデルを明示的にキャッシュに保存することが重要です。これにより、ユーザーがアプリを再読み込みした後にモデルデータをすぐに利用できるようになります。

これを実現するために使用できる手法はいくつかあります。次のコードサンプルでは、各モデルファイルがメモリ内の blob という名前の Blob オブジェクトに保存されていると想定しています。

パフォーマンスを把握するために、各コードサンプルには performance.mark() メソッドと performance.measure() メソッドがアノテーションされています。これらの測定値はデバイスに依存し、一般化できません。

Chrome DevTools の [Application] > [Storage] で、IndexedDB、キャッシュ ストレージ、ファイル システムのセグメントを含む使用状況図を確認します。各セグメントで 1,354 メガバイトのデータが消費され、合計で 4,063 メガバイトになります。

ブラウザで AI モデルをキャッシュに保存するには、Cache APIOrigin Private File System APIIndexedDB API のいずれかを使用できます。一般に推奨されるのは Cache API を使用することですが、このガイドではすべてのオプションの長所と短所について説明します。

Cache API

Cache API は、長期間のメモリにキャッシュに保存される Request オブジェクトペアと Response オブジェクトペア用の永続ストレージを提供します。Service Workers の仕様で定義されていますが、この API はメインスレッドまたは通常のワーカーから使用できます。サービス ワーカーのコンテキスト外で使用するには、Request オブジェクトではなく合成 URL とペア設定された合成 Response オブジェクトを使用して Cache.put() メソッドを呼び出します。

このガイドでは、インメモリ blob を前提としています。キャッシュキーとして偽の URL を使用し、blob に基づいて合成 Response を使用します。モデルを直接ダウンロードする場合は、fetch() リクエストから取得した Response を使用します。

たとえば、Cache API を使用してモデルファイルを保存して復元する方法は次のとおりです。

const storeFileInSWCache = async (blob) => {
  try {
    performance.mark('start-sw-cache-cache');
    const modelCache = await caches.open('models');
    await modelCache.put('model.bin', new Response(blob));
    performance.mark('end-sw-cache-cache');

    const mark = performance.measure(
      'sw-cache-cache',
      'start-sw-cache-cache',
      'end-sw-cache-cache'
    );
    console.log('Model file cached in sw-cache.', mark.name, mark.duration.toFixed(2));
  } catch (err) {
    console.error(err.name, err.message);
  }
};

const restoreFileFromSWCache = async () => {
  try {
    performance.mark('start-sw-cache-restore');
    const modelCache = await caches.open('models');
    const response = await modelCache.match('model.bin');
    if (!response) {
      throw new Error(`File model.bin not found in sw-cache.`);
    }
    const file = await response.blob();
    performance.mark('end-sw-cache-restore');
    const mark = performance.measure(
      'sw-cache-restore',
      'start-sw-cache-restore',
      'end-sw-cache-restore'
    );
    console.log(mark.name, mark.duration.toFixed(2));
    console.log('Cached model file found in sw-cache.');
    return file;
  } catch (err) {    
    throw err;
  }
};

Origin Private File System API

オリジン プライベート ファイル システム(OPFS)は、ストレージ エンドポイントの比較的新しい標準です。ページのオリジンに限定されるため、通常のファイル システムとは異なり、ユーザーには表示されません。パフォーマンスを重視して高度に最適化された特別なファイルへのアクセスを提供し、そのコンテンツへの書き込みアクセスを提供します。

たとえば、OPFS にモデルファイルを保存して復元する方法は次のとおりです。

const storeFileInOPFS = async (blob) => {
  try {
    performance.mark('start-opfs-cache');
    const root = await navigator.storage.getDirectory();
    const handle = await root.getFileHandle('model.bin', { create: true });
    const writable = await handle.createWritable();
    await blob.stream().pipeTo(writable);
    performance.mark('end-opfs-cache');
    const mark = performance.measure(
      'opfs-cache',
      'start-opfs-cache',
      'end-opfs-cache'
    );
    console.log('Model file cached in OPFS.', mark.name, mark.duration.toFixed(2));
  } catch (err) {
    console.error(err.name, err.message);
  }
};

const restoreFileFromOPFS = async () => {
  try {
    performance.mark('start-opfs-restore');
    const root = await navigator.storage.getDirectory();
    const handle = await root.getFileHandle('model.bin');
    const file = await handle.getFile();
    performance.mark('end-opfs-restore');
    const mark = performance.measure(
      'opfs-restore',
      'start-opfs-restore',
      'end-opfs-restore'
    );
    console.log('Cached model file found in OPFS.', mark.name, mark.duration.toFixed(2));
    return file;
  } catch (err) {    
    throw err;
  }
};

IndexedDB API

IndexedDB は、任意のデータをブラウザに永続的に保存するための確立された標準です。IndexedDB は API が複雑なことで悪名高いですが、idb-keyval などのラッパー ライブラリを使用すると、IndexedDB を従来のキーバリュー ストアのように扱うことができます。

次に例を示します。

import { get, set } from 'https://cdn.jsdelivr.net/npm/idb-keyval@latest/+esm';

const storeFileInIDB = async (blob) => {
  try {
    performance.mark('start-idb-cache');
    await set('model.bin', blob);
    performance.mark('end-idb-cache');
    const mark = performance.measure(
      'idb-cache',
      'start-idb-cache',
      'end-idb-cache'
    );
    console.log('Model file cached in IDB.', mark.name, mark.duration.toFixed(2));
  } catch (err) {
    console.error(err.name, err.message);
  }
};

const restoreFileFromIDB = async () => {
  try {
    performance.mark('start-idb-restore');
    const file = await get('model.bin');
    if (!file) {
      throw new Error('File model.bin not found in IDB.');
    }
    performance.mark('end-idb-restore');
    const mark = performance.measure(
      'idb-restore',
      'start-idb-restore',
      'end-idb-restore'
    );
    console.log('Cached model file found in IDB.', mark.name, mark.duration.toFixed(2));
    return file;
  } catch (err) {    
    throw err;
  }
};

ストレージを永続としてマークする

これらのキャッシュ保存方法のいずれかの末尾で navigator.storage.persist() を呼び出して、永続ストレージの使用許可をリクエストします。このメソッドは、権限が付与されている場合は true、そうでない場合は false に解決する Promise を返します。ブラウザ固有のルールに応じて、ブラウザはリクエストに応じる場合もあれば、応じない場合もあります

if ('storage' in navigator && 'persist' in navigator.storage) {
  try {
    const persistent = await navigator.storage.persist();
    if (persistent) {
      console.log("Storage will not be cleared except by explicit user action.");
      return;
    }
    console.log("Storage may be cleared under storage pressure.");  
  } catch (err) {
    console.error(err.name, err.message);
  }
}

特別なケース: ハードディスクでモデルを使用する

ブラウザ ストレージの代わりに、ユーザーのハードディスクから AI モデルを直接参照できます。この手法は、研究に重点を置いたアプリで特定のモデルをブラウザで実行する可能性を示すのに役立ちます。また、アーティストが専門的なクリエイティビティ アプリでセルフ トレーニング モデルを使用できるようにすることもできます。

File System Access API

File System Access API を使用すると、ハードディスクからファイルを開き、IndexedDB に保持できる FileSystemFileHandle を取得できます。

このパターンでは、ユーザーはモデルファイルへのアクセス権を 1 回だけ付与する必要があります。永続的な権限により、ユーザーはファイルへのアクセスを恒久的に許可できます。アプリを再読み込みし、マウスクリックなどの必要なユーザー操作を行った後、IndexedDB からハードディスク上のファイルにアクセスして FileSystemFileHandle を復元できます。

ファイル アクセス権はクエリされ、必要に応じてリクエストされるため、今後の再読み込みがシームレスになります。次の例は、ハードディスクからファイルのハンドルを取得し、ハンドルを保存して復元する方法を示しています。

import { fileOpen } from 'https://cdn.jsdelivr.net/npm/browser-fs-access@latest/dist/index.modern.js';
import { get, set } from 'https://cdn.jsdelivr.net/npm/idb-keyval@latest/+esm';

button.addEventListener('click', async () => {
  try {
    const file = await fileOpen({
      extensions: ['.bin'],
      mimeTypes: ['application/octet-stream'],
      description: 'AI model files',
    });
    if (file.handle) {
      // It's an asynchronous method, but no need to await it.
      storeFileHandleInIDB(file.handle);
    }
    return file;
  } catch (err) {
    if (err.name !== 'AbortError') {
      console.error(err.name, err.message);
    }
  }
});

const storeFileHandleInIDB = async (handle) => {
  try {
    performance.mark('start-file-handle-cache');
    await set('model.bin.handle', handle);
    performance.mark('end-file-handle-cache');
    const mark = performance.measure(
      'file-handle-cache',
      'start-file-handle-cache',
      'end-file-handle-cache'
    );
    console.log('Model file handle cached in IDB.', mark.name, mark.duration.toFixed(2));
  } catch (err) {
    console.error(err.name, err.message);
  }
};

const restoreFileFromFileHandle = async () => {
  try {
    performance.mark('start-file-handle-restore');
    const handle = await get('model.bin.handle');
    if (!handle) {
      throw new Error('File handle model.bin.handle not found in IDB.');
    }
    if ((await handle.queryPermission()) !== 'granted') {
      const decision = await handle.requestPermission();
      if (decision === 'denied' || decision === 'prompt') {
        throw new Error(Access to file model.bin.handle not granted.');
      }
    }
    const file = await handle.getFile();
    performance.mark('end-file-handle-restore');
    const mark = performance.measure(
      'file-handle-restore',
      'start-file-handle-restore',
      'end-file-handle-restore'
    );
    console.log('Cached model file handle found in IDB.', mark.name, mark.duration.toFixed(2));
    return file;
  } catch (err) {    
    throw err;
  }
};

これらの方法は相互排他的ではありません。ブラウザでモデルを明示的にキャッシュし、ユーザーのハードディスクからモデルを使用する場合もあります。

デモ

3 つの通常のケース ストレージ方法と、MediaPipe LLM デモで実装されているハードディスク メソッドを確認できます。

ボーナス: サイズの大きいファイルをチャンクでダウンロードする

インターネットから大きな AI モデルをダウンロードする必要がある場合は、ダウンロードを個別のチャンクに並列化し、クライアントで再びステッチします。

コードで使用できるヘルパー関数を次に示します。url を渡すだけで、chunkSize(デフォルト: 5 MB)、maxParallelRequests(デフォルト: 6)、progressCallback 関数(downloadedBytes と合計 fileSize を報告)、AbortSignal シグナルの signal はすべて省略可能です。

次の関数をプロジェクトにコピーするか、npm から fetch-in-chunks パッケージをインストールします。

async function fetchInChunks(
  url,
  chunkSize = 5 * 1024 * 1024,
  maxParallelRequests = 6,
  progressCallback = null,
  signal = null
) {
  // Helper function to get the size of the remote file using a HEAD request
  async function getFileSize(url, signal) {
    const response = await fetch(url, { method: 'HEAD', signal });
    if (!response.ok) {
      throw new Error('Failed to fetch the file size');
    }
    const contentLength = response.headers.get('content-length');
    if (!contentLength) {
      throw new Error('Content-Length header is missing');
    }
    return parseInt(contentLength, 10);
  }

  // Helper function to fetch a chunk of the file
  async function fetchChunk(url, start, end, signal) {
    const response = await fetch(url, {
      headers: { Range: `bytes=${start}-${end}` },
      signal,
    });
    if (!response.ok && response.status !== 206) {
      throw new Error('Failed to fetch chunk');
    }
    return await response.arrayBuffer();
  }

  // Helper function to download chunks with parallelism
  async function downloadChunks(
    url,
    fileSize,
    chunkSize,
    maxParallelRequests,
    progressCallback,
    signal
  ) {
    let chunks = [];
    let queue = [];
    let start = 0;
    let downloadedBytes = 0;

    // Function to process the queue
    async function processQueue() {
      while (start < fileSize) {
        if (queue.length < maxParallelRequests) {
          let end = Math.min(start + chunkSize - 1, fileSize - 1);
          let promise = fetchChunk(url, start, end, signal)
            .then((chunk) => {
              chunks.push({ start, chunk });
              downloadedBytes += chunk.byteLength;

              // Update progress if callback is provided
              if (progressCallback) {
                progressCallback(downloadedBytes, fileSize);
              }

              // Remove this promise from the queue when it resolves
              queue = queue.filter((p) => p !== promise);
            })
            .catch((err) => {              
              throw err;              
            });
          queue.push(promise);
          start += chunkSize;
        }
        // Wait for at least one promise to resolve before continuing
        if (queue.length >= maxParallelRequests) {
          await Promise.race(queue);
        }
      }

      // Wait for all remaining promises to resolve
      await Promise.all(queue);
    }

    await processQueue();

    return chunks.sort((a, b) => a.start - b.start).map((chunk) => chunk.chunk);
  }

  // Get the file size
  const fileSize = await getFileSize(url, signal);

  // Download the file in chunks
  const chunks = await downloadChunks(
    url,
    fileSize,
    chunkSize,
    maxParallelRequests,
    progressCallback,
    signal
  );

  // Stitch the chunks together
  const blob = new Blob(chunks);

  return blob;
}

export default fetchInChunks;

自分に合った方法を選択する

このガイドでは、ブラウザで AI モデルを効果的にキャッシュに保存するさまざまな方法について説明しました。これは、アプリのユーザー エクスペリエンスとパフォーマンスを向上させるうえで重要なタスクです。Chrome ストレージ チームは、パフォーマンスを最適化するために Cache API を使用することをおすすめします。これにより、AI モデルへの迅速なアクセスが確保され、読み込み時間が短縮され、応答性が向上します。

OPFS と IndexedDB は、あまり使用できないオプションです。OPFS API と IndexedDB API では、データを保存する前にデータをシリアル化する必要があります。また、IndexedDB ではデータを取得するときに逆シリアル化を行う必要があるため、大規模なモデルを保存するには適していません。

ニッチなアプリケーションの場合、File System Access API を使用すると、ユーザーのデバイス上のファイルに直接アクセスできます。これは、独自の AI モデルを管理するユーザーに最適です。

AI モデルを保護する必要がある場合は、サーバーに保持します。クライアントに保存されたデータは、DevTools または OFPS DevTools 拡張機能を使用して、キャッシュと IndexedDB の両方から簡単に抽出できます。これらのストレージ API は、本質的にセキュリティが同等です。暗号化されたバージョンのモデルを保存することもできますが、その場合は復号鍵をクライアントに渡す必要があります。この鍵はインターセプトされる可能性があります。つまり、悪意のある人物がモデルを盗もうとするのは少し難しくなりますが、不可能ではありません。

アプリの要件、ターゲット オーディエンスの動作、使用される AI モデルの特性に合わせて、キャッシュ戦略を選択することをおすすめします。これにより、さまざまなネットワーク条件とシステム制約下で、アプリケーションの応答性と堅牢性が確保されます。


謝辞

ジョシュア ベル、ライリー グラント、エヴァン スタッド、ネイサン メモット、オースティン サリバン、エティエンヌ ノエル、アンドレ バンダラ、アレクサンドラ クレッパー、フランソワ ボーフォール、ポール キンラン、レイチェル アンドリューが確認しました。