使用 TensorFlow.js 產生大小最佳化的瀏覽器套件

總覽

TensorFlow.js 3.0 支援建構大小最佳化、以生產環境為導向的瀏覽器套件。換句話說,我們希望讓您更輕鬆地將更少的 JavaScript 傳送到瀏覽器。

此功能適用於生產環境使用案例的使用者,他們特別希望減少有效負載的位元組 (因此願意投入精力來達成此目標)。若要使用此功能,您應熟悉 ES 模組、JavaScript 套件工具 (例如 webpackrollup) 以及 樹狀結構修剪/無效程式碼消除等概念。

本教學課程示範如何建立自訂 tensorflow.js 模組,該模組可與套件工具搭配使用,為使用 tensorflow.js 的程式產生大小最佳化的建構。

詞彙

在本文件中,我們將使用一些關鍵詞彙

ES 模組 - 標準 JavaScript 模組系統。在 ES6/ES2015 中導入。可透過使用 importexport 陳述式來識別。

套件 - 取得一組 JavaScript 資產,並將其分組/捆綁成一個或多個可在瀏覽器中使用的 JavaScript 資產。這是通常產生要提供給瀏覽器的最終資產的步驟。應用程式通常會直接從轉譯的程式庫來源執行自己的套件作業常見的套件工具包括 rollupwebpack。套件作業的最終結果稱為套件 (或有時稱為區塊,如果它分成多個部分)

樹狀結構修剪/無效程式碼消除 - 移除最終撰寫的應用程式未使用的程式碼。這是在套件作業期間完成的,通常在最小化步驟中完成。

運算 (Ops) - 對一個或多個張量執行數學運算,產生一個或多個張量作為輸出。運算是「高階」程式碼,可以使用其他運算來定義其邏輯。

核心 - 運算的特定實作,與特定硬體功能相關聯。核心是「低階」且後端特定的。某些運算具有從運算到核心的一對一對應,而其他運算則使用多個核心。

範圍和使用案例

僅推論圖形模型

我們從使用者那裡聽說的與此相關的主要使用案例,以及我們在此版本中支援的是使用 TensorFlow.js 圖形模型 進行推論。如果您使用 TensorFlow.js 層模型,可以使用 tfjs-converter 將其轉換為圖形模型格式。圖形模型格式對於推論使用案例更有效率。

使用 tfjs-core 進行低階張量操作

我們支援的另一個使用案例是直接使用 @tensorflow/tjfs-core 套件進行較低階張量操作的程式。

我們的自訂建構方法

我們在設計此功能時的核心原則包括以下內容

  • 盡可能充分利用 JavaScript 模組系統 (ESM),並允許 TensorFlow.js 的使用者也這樣做。
  • 盡可能使 TensorFlow.js 可透過現有的套件工具 (例如 webpack、rollup 等) 進行樹狀結構修剪。這讓使用者能夠充分利用這些套件工具的所有功能,包括程式碼分割等功能。
  • 盡可能讓對套件大小不那麼敏感的使用者更容易使用。這確實表示生產環境建構將需要更多努力,因為我們程式庫中的許多預設值都支援易用性,而不是大小最佳化的建構。

我們工作流程的主要目標是為 TensorFlow.js 產生自訂 JavaScript 模組,其中僅包含我們嘗試最佳化之程式所需的功能。我們仰賴現有的套件工具來執行實際的最佳化。

雖然我們主要仰賴 JavaScript 模組系統,但我們也提供自訂 CLI 工具來處理使用者介面程式碼中不易透過模組系統指定的部分。以下是兩個範例:

  • 儲存在 model.json 檔案中的模型規格
  • 我們使用的運算至後端特定核心調度系統。

這使得產生自訂 tfjs 建構比僅將套件工具指向一般的 @tensorflow/tfjs 套件更複雜一些。

如何建立大小最佳化的自訂套件

步驟 1:判斷您的程式正在使用哪些核心

此步驟可讓我們判斷您執行的任何模型或預先/後處理程式碼 (在您選取的後端下) 使用的所有核心。

使用 tf.profile 執行應用程式中使用 tensorflow.js 的部分,並取得核心。它看起來會像這樣

const profileInfo = await tf.profile(() => {
  // You must profile all uses of tf symbols.
  runAllMyTfjsCode();
});

const kernelNames = profileInfo.kernelNames
console.log(kernelNames);

將該核心清單複製到剪貼簿,以進行下一步。

您需要使用與您要在自訂套件中使用的相同後端來分析程式碼。

如果您的模型變更或您的預先/後處理程式碼變更,您將需要重複此步驟。

步驟 2. 撰寫自訂 tfjs 模組的組態檔

以下是組態檔範例。

它看起來像這樣

{
  "kernels": ["Reshape", "_FusedMatMul", "Identity"],
  "backends": [
      "cpu"
  ],
  "models": [
      "./model/model.json"
  ],
  "outputPath": "./custom_tfjs",
  "forwardModeOnly": true
}
  • kernels:要包含在套件中的核心清單。從步驟 1 的輸出複製此清單。
  • backends:您要包含的後端清單。有效選項包括 "cpu"、"webgl" 和 “wasm”。
  • models:應用程式中載入之模型的 model.json 檔案清單。如果您的程式未使用 tfjs_converter 來載入圖形模型,則可以為空。
  • outputPath:放置產生模組的資料夾路徑。
  • forwardModeOnly:如果您要包含先前列出的核心的梯度,請將此設定為 false。

步驟 3. 產生自訂 tfjs 模組

以組態檔作為引數執行自訂建構工具。您需要安裝 @tensorflow/tfjs 套件才能存取此工具。

npx tfjs-custom-module  --config custom_tfjs_config.json

這將在 outputPath 建立一個資料夾,其中包含一些新檔案。

步驟 4. 設定您的套件工具,將 tfjs 別名設定為新的自訂模組。

在 webpack 和 rollup 等套件工具中,我們可以將現有對 tfjs 模組的參考別名設定為指向我們新產生的自訂 tfjs 模組。有三個模組需要設定別名,才能最大程度地節省套件大小。

以下是在 webpack 中的範例程式碼片段 (完整範例在此)

...

config.resolve = {
  alias: {
    '@tensorflow/tfjs$':
        path.resolve(__dirname, './custom_tfjs/custom_tfjs.js'),
    '@tensorflow/tfjs-core$': path.resolve(
        __dirname, './custom_tfjs/custom_tfjs_core.js'),
    '@tensorflow/tfjs-core/dist/ops/ops_for_converter': path.resolve(
        __dirname, './custom_tfjs/custom_ops_for_converter.js'),
  }
}

...

以下是 rollup 的等效程式碼片段 (完整範例在此)

import alias from '@rollup/plugin-alias';

...

alias({
  entries: [
    {
      find: /@tensorflow\/tfjs$/,
      replacement: path.resolve(__dirname, './custom_tfjs/custom_tfjs.js'),
    },
    {
      find: /@tensorflow\/tfjs-core$/,
      replacement: path.resolve(__dirname, './custom_tfjs/custom_tfjs_core.js'),
    },
    {
      find: '@tensorflow/tfjs-core/dist/ops/ops_for_converter',
      replacement: path.resolve(__dirname, './custom_tfjs/custom_ops_for_converter.js'),
    },
  ],
}));

...

如果您的套件工具不支援模組別名,您需要變更您的 import 陳述式,以從步驟 3 中建立的產生 custom_tfjs.js 匯入 tensorflow.js。運算定義將不會透過樹狀結構修剪掉,但核心仍然會透過樹狀結構修剪掉。一般而言,樹狀結構修剪核心是最終套件大小的最大節省來源。

如果您僅使用 @tensoflow/tfjs-core 套件,則只需要別名設定該套件即可。

步驟 5. 建立您的套件

執行您的套件工具 (例如 webpackrollup) 以產生您的套件。套件的大小應小於在沒有模組別名設定的情況下執行套件工具的大小。您也可以使用視覺化工具 (例如 這個工具) 來查看最終套件中包含的內容。

步驟 6. 測試您的應用程式

請務必測試您的應用程式是否如預期般運作!