在浏览器中缓存 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,其中已选择三个对象:两只长颈鹿和一颗月亮。

为了在以后加快应用的启动速度,您应该在设备上显式缓存模型数据,而不是依赖于隐式 HTTP 浏览器缓存。

虽然本指南使用 gemma-2b-it-gpu-int4.bin model 来创建聊天机器人,但该方法还可以进行泛化,以适应设备上的其他模型和其他用例。将应用关联到模型的最常用方法是将模型与其余应用资源一起提供。优化投放至关重要

配置合适的缓存标头

如果您从服务器提供 AI 模型,请务必配置正确的 Cache-Control 标头。以下示例展示了可靠的默认设置,您可以根据应用的需求来构建该设置。

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

AI 模型的每个发布版本都是静态资源。对于从不更改的内容,应在请求网址中为其提供较长的 max-age缓存无效化。如果您确实需要更新模型,则必须为其提供一个新网址

当用户重新加载页面时,即使服务器知道内容处于稳定状态,客户端也会发送重新验证请求。immutable 指令明确指出不需要重新验证,因为内容不会更改。immutable 指令未得到浏览器和中间缓存或代理服务器的广泛支持,但通过将其与公认的 max-age 指令结合使用,可以确保实现最大的兼容性。public 响应指令指示响应可以存储在共享缓存中。

在请求 AI 模型时,Chrome 开发者工具会显示 Hugging Face 发送的生产 Cache-Control 标头。 (来源

在客户端缓存 AI 模型

在提供 AI 模型时,请务必在浏览器中明确缓存该模型。这可确保模型数据在用户重新加载应用后随时可用。

您可以使用许多技术来实现这一目标。对于以下代码示例,假设每个模型文件都存储在内存中名为 blobBlob 对象中。

为了解性能,每个代码示例都带有 performance.mark()performance.measure() 方法注解。这些措施取决于设备,无法泛化。

在 Chrome 开发者工具的 Application > Storage 中,查看使用情况图,其中包含 IndexedDB、Cache 存储和文件系统的细分。 所显示的每个细分消耗了 1354 MB 的数据,总计 4063 MB。

您可以选择使用以下 API 之一在浏览器中缓存 AI 模型:Cache APIOrigin Private File System APIIndexedDB API一般建议使用 Cache API,但本指南将讨论所有选项的优缺点。

Cache API

Cache API 为在长效内存中缓存的 RequestResponse 对象对提供永久性存储空间。尽管您可从主线程或常规工作器使用该 API,但在 Service Worker 规范中进行了定义。如需在 Service Worker 上下文之外使用它,请使用合成的 Response 对象调用 Cache.put() 方法,该对象与合成网址而不是 Request 对象配对。

本指南假定使用内存中 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;
  }
};

源私有文件系统 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;
  }
};

索引型数据库 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 的 promise,否则返回 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 演示中看到全部三种常规 case 存储方法和硬盘方法。

额外提示:分块下载大型文件

如果您需要从互联网下载一个大型 AI 模型,请将其并行下载为单独的数据块,然后在客户端上再次拼接在一起。

以下是可以在代码中使用的辅助函数。您只需向其传递 urlchunkSize(默认值:5MB)、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 还需要在检索数据时对其进行反序列化,这使其成为最适合存储大型模型的位置。

对于小众应用,File System Access API 可让您直接访问用户设备上的文件,非常适合自行管理 AI 模型的用户。

如果您需要保护 AI 模型,请将其保留在服务器上。将数据存储在客户端后,您可以使用开发者工具或 OFPS 开发者工具扩展程序从 Cache 和 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 审核。