在瀏覽器中快取 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 的示範已存在,但以下是一些在瀏覽器中執行的其他模型的實際工作環境級範例:

在 Adobe Photoshop 線上版中開啟 AI 技術輔助的物件選取工具,並選取三個物件:兩隻長頸鹿和月亮。

為加快日後應用程式的啟動速度,您應明確快取裝置端的模型資料,而非依賴隱含的 HTTP 瀏覽器快取。

雖然本指南使用 gemma-2b-it-gpu-int4.bin model 建立對話方塊,但這項做法可概略套用至其他裝置模型和用途。將應用程式連結至模型最常見的方式,就是與其他應用程式資源一併提供模型。因此,請務必將提交作業最佳化。

設定正確的快取標頭

如果您從伺服器提供 AI 模型,請務必設定正確的 Cache-Control 標頭。以下範例顯示可靠的預設設定,您可以根據應用程式需求進行建構。

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

每個 AI 模型的發布版本都是靜態資源。從未變更的內容應在要求網址中,搭配使用長 max-age快取破壞。如果確實需要更新模型,請務必為其提供新網址

使用者重新載入網頁時,即使伺服器知道內容穩定,用戶端仍會傳送重新驗證要求。immutable 指令明確指出內容不會變更,因此不需要重新驗證。瀏覽器和中介快取或 Proxy 伺服器不廣泛支援 immutable 指令,但如果搭配普遍認可的 max-age 指令,就能確保最大相容性。public 回應指令表示回應可儲存在共用快取中。

Chrome 開發人員工具會顯示 Hugging Face 在要求 AI 模型時傳送的正式版 Cache-Control 標頭。(來源)

在用戶端快取 AI 模型

提供 AI 模型時,請務必在瀏覽器中明確快取模型。這可確保使用者重新載入應用程式後,模型資料隨時可用。

您可以使用多種技巧來達成這項目標。針對下列程式碼範例,假設每個模型檔案都儲存在記憶體中名為 blobBlob 物件中。

為瞭解效能,每個程式碼範例都會加上 performance.mark()performance.measure() 方法的註解。這些措施會因裝置而異,無法一體適用。

在 Chrome 開發人員工具的「應用程式」 >「儲存空間」中,查看包含 IndexedDB、快取儲存空間和檔案系統的使用率圖表。每個區段的資料使用量為 1354 MB,總計為 4063 MB。

您可以選擇使用下列任一 API,在瀏覽器中快取 AI 模型:Cache APIOrigin Private File System APIIndexedDB API一般建議使用 Cache API,但本指南會討論所有選項的優缺點。

Cache API

Cache API 可為在長效記憶體中快取的 RequestResponse 物件組合提供持續性儲存空間。雖然此 API 是在 Service Workers 規格中定義,但您可以從主執行緒或一般 worker 使用此 API。如要在服務工作者背景處理程序外使用,請使用合成 Response 物件搭配合成網址 (而非 Request 物件) 呼叫 Cache.put() 方法。

本指南假設記憶體內有 blob。使用假網址做為快取金鑰,並根據 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

Origin Private File System (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 是一種成熟的標準,可讓您以持久方式在瀏覽器中儲存任意資料。它以 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。瀏覽器可能會或不會接受要求,這取決於瀏覽器專屬規則。

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

使用這個模式時,使用者只需授予一次模型檔案存取權。由於有持續性權限,使用者可以選擇永久授予檔案存取權。重新載入應用程式並執行必要的使用者手勢 (例如滑鼠點擊) 後,即可透過 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;
  }
};

這些方法並非互斥。在某些情況下,您可能會在瀏覽器中明確快取模型,並使用使用者硬碟中的模型。

示範

您可以在 MediaPipe LLM 示範中查看所有三種一般儲存空間方法,以及實作硬碟方法。

額外說明:分塊下載大型檔案

如果您需要從網際網路下載大型 AI 模型,請將下載作業並行化為個別區塊,然後在用戶端上重新拼接。

以下是可在程式碼中使用的輔助函式。您只需傳遞 urlchunkSize (預設值: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 和 IndexedDB API 需要先將資料序列化,才能儲存資料。IndexedDB 在擷取資料時也需要將資料反序列化,因此不適合用於儲存大型模型。

對於特定應用程式,檔案系統存取 API 可讓您直接存取使用者裝置上的檔案,非常適合自行管理 AI 模型的使用者。

如果您需要保護 AI 模型,請將模型保留在伺服器上。一旦儲存到用戶端,只要使用開發人員工具或 OFPS 開發人員工具擴充功能,就能輕鬆從快取和 IndexedDB 中擷取資料。這些儲存空間 API 在安全性方面本質上是相同的。您可能會想儲存加密版本的模型,但接著您需要將解密金鑰傳送至可能遭到攔截的用戶端。也就是說,惡意人士要竊取模型的難度會稍微提高,但並非不可能。

建議您選擇符合應用程式需求、目標對象行為和所用 AI 模型特性的快取策略。這可確保應用程式在各種網路條件和系統限制下,都能保持回應性和穩定性。


特別銘謝

本文由 Joshua Bell、Reilly Grant、Evan Stade、Nathan Memmott、Austin Sullivan、Etienne Noël、André Bandarra、Alexandra Klepper、François Beaufort、Paul Kinlan 和 Rachel Andrew 共同審查。