處理大型模型

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

故障排除

本節正在建設中。