处理大型模型

ONNX 模型的大小因模型复杂度和参数数量而异。它们可以小到几 KB,也可以大到几 GB。虽然 ONNX Runtime Web 旨在在浏览器中运行所有模型,但在处理大型模型时仍需注意一些事项。

目录

平台限制

在浏览器中处理大型模型时,您应该注意一些平台限制

ArrayBuffer 的最大大小

尽管 JavaScript 中的 ArrayBuffer 大小没有硬性限制,但每个浏览器都有其自身的限制。例如,Chrome 中 ArrayBuffer 的最大大小为 0x7fe00000 字节(约 2GB)。在使用 fetch API 加载大型模型时请务必小心,因为如果在大型文件上调用 response.arrayBuffer() 可能会失败。

ONNX Runtime Web 通过使用 new WebAssembly.Memory() 创建大于 2GB 的 ArrayBuffer,从而绕过此限制。但是,通过 new WebAssembly.Memory() 创建的 ArrayBuffer 实例不可转移,因此无法与 Proxy 功能配合使用。

Protobuf 文件大小限制

ONNX 模型以 protobuf 格式序列化。protobuf 文件的最大大小为 2GB。如果 ONNX 模型大于 2GB,通常会生成外部数据。有关更多详细信息,请参阅外部数据

WebAssembly 内存限制

WebAssembly 的内存限制为 4GB。由于 32 位寻址,这是 WebAssembly 模块可以访问的最大内存量。目前,ONNX Runtime Web 无法运行大于 4GB 的模型。我们将来可能会通过使用 WASM64 或直接加载 GPU 权重来支持它。

缓存模型

为了避免每次刷新页面时都加载模型,您可以使用 Cache API 或 Origin Private File System 来缓存模型。这样,模型就可以从缓存中加载,而不是每次都从服务器获取。

有关更多详细信息,请参阅 Cache APIOrigin private file system

外部数据

当您处理大型 ONNX 模型时,它通常会伴随外部数据一起生成。由于 protobuf 文件大小限制,大于 2GB 的 ONNX 模型必须使用外部数据。外部数据是一个或多个独立文件,通常由 ONNX 导出器生成。外部数据通常与 ONNX 模型文件放在同一个目录中。

ONNX Runtime 支持加载带外部数据的模型。这在 C/C++/Python API 中无需额外步骤即可自动完成,因为这些语言绑定的 ONNX Runtime 可以访问文件系统。然而,在浏览器中,JavaScript 代码无法直接访问文件系统。因此,您需要额外一步将外部数据信息传递给 ONNX Runtime Web。

外部数据工作原理

在深入细节之前,我们首先了解一下 ONNX Runtime 中外部数据的工作原理。此信息很重要,否则步骤可能会令人困惑。

ONNX 模型在技术上是一个 protobuf 文件。protobuf 文件包含模型图和权重。ONNX 规范允许将权重存储在 protobuf 文件中或外部数据文件中。当权重存储在 protobuf 文件中时,它完全包含在 protobuf 文件中。当权重存储在外部数据文件中时,protobuf 文件包含该特定权重的以下信息:

  • “location”(一个字符串),指定外部数据文件的相对文件路径
  • “offset”(一个整数),指定外部数据文件中权重开始的字节偏移量
  • “length”(一个整数),指定权重的字节长度

“location” 通常由 ONNX 导出器确定。例如,导出器可能将模型文件输出为 model_a.onnx,将外部数据文件输出为 model_a.data 放在同一目录中。模型中的某些权重存储在 model_a.data 文件中,因此这些权重的“location”被设置为 ./model_a.data。此信息存储在 model_a.onnx 文件中。

这解释了为什么在原生平台上,确保绝不重命名外部数据文件很重要。如果您重命名外部数据文件,protobuf 文件中的“location”将与实际文件名不匹配,ONNX Runtime 将无法加载模型。

对于 ONNX Runtime Web,您始终需要将外部数据信息传递给 ONNX Runtime Web。理解 protobuf 文件中定义的“location”与实际外部文件路径是两个不同的概念非常重要。这两个不同的概念将在下一节的 JavaScript 代码中分别表示为“path”和“data”。

在 ONNX Runtime Web 中加载带外部数据的模型

我们使用一个示例来说明如何在浏览器中加载带外部数据的 ONNX 模型。假设我们有一个 ONNX 模型 model_a.onnx 和一个外部数据文件 model_a.data。以下代码展示了如何在浏览器中加载带外部数据的模型:

const modelUrl = 'https://example.com/path/model_a.onnx';
const externalDataUrl = 'https://example.com/path/model_a.data';

const mySession = await ort.InferenceSession.create(modelUrl, {
    ...,
    externalData: [
        {
            path: './model_a.data',
            data: externalDataUrl
        }
    ]
});

在上面的代码中,我们将外部数据信息传递给 InferenceSession.create() 方法。externalData 是一个对象数组,每个对象代表一个外部数据文件。该对象有两个属性:

  • path(一个字符串),应与 protobuf 文件中权重的“location”信息匹配
  • data(一个字符串),指定外部数据文件。它可以是 URL、Blob 或 Uint8Array。

当您将模型和外部数据存储在 IndexedDB 中时,您可以从 IndexedDB 加载带外部数据的模型。以下代码展示了如何从 IndexedDB 加载带外部数据的模型:

// assume loadFromIndexedDB() is a function implemented by your app that loads the data from the IndexedDB

// Load the model and external data from the IndexedDB
const modelBlob = await loadFromIndexedDB('model_a.onnx');
const externalDataBlob = await loadFromIndexedDB('model_a.data');

const mySession = await ort.InferenceSession.create(modelBlob, {
    ...,
    externalData: [
        {
            path: './model_a.data',
            data: externalDataBlob
        }
    ]
});

有关更多详细信息,请参阅 ONNX External Data

故障排除

本节正在建设中。