TensorFlow.js を使用してサイズが最適化されたブラウザ バンドルを生成する

概要

TensorFlow.js 3.0 は、サイズが最適化された実稼働指向のブラウザ バンドルの構築をサポートします。別の言い方をすれば、ブラウザに送信する JavaScript の量を減らしやすくしたいと考えています。

この機能は、ペイロードからバイト数を削減することで特に恩恵を受ける (したがって、これを達成するために喜んで努力する) 運用ユースケースを持つユーザーを対象としています。この機能を使用するには、 ES モジュールwebpackrollupなどの JavaScript バンドル ツール、およびツリーシェイキング/デッドコード除去などの概念に精通している必要があります。

このチュートリアルでは、tensorflow.js を使用してプログラムのサイズが最適化されたビルドを生成するためにバンドラーで使用できるカスタム tensorflow.js モジュールを作成する方法を示します。

用語

このドキュメントの文脈では、いくつかの重要な用語が使用されます。

ES モジュール-標準の JavaScript モジュール システム。 ES6/ES2015で導入されました。 importおよびexportステートメントを使用して識別できます。

バンドル- 一連の JavaScript アセットを取得し、ブラウザーで使用できる 1 つ以上の JavaScript アセットにグループ化/バンドルすること。これは通常、ブラウザに提供される最終的なアセットを生成するステップです。通常、アプリケーションは、トランスパイルされたライブラリ ソースから直接独自のバンドルを実行します一般的なバンドラーには、 rollupおよびwebpackが含まれます。バンドルの最終結果はバンドルと呼ばれます (複数の部分に分割されている場合はチャンクとして知られることもあります)。

Tree-Shaking / Dead Code Elimination - 最終的に作成されたアプリケーションで使用されないコードを削除します。これはバンドル中に、通常は縮小ステップで行われます。

演算 (Ops) - 1 つまたは複数のテンソルを出力として生成する、1 つまたは複数のテンソルに対する数学的演算。演算は「高レベル」コードであり、他の演算を使用してロジックを定義できます。

カーネル- 特定のハードウェア機能に関連付けられたオプの特定の実装。カーネルは「低レベル」でバックエンド固有です。一部の操作は操作からカーネルへ 1 対 1 のマッピングを持ちますが、他の操作は複数のカーネルを使用します。

範囲と使用例

推論のみのグラフモデル

これに関連してユーザーから聞き、このリリースでサポートしている主な使用例は、 TensorFlow.js グラフ モデルを使用して推論を行う使用例です。 TensorFlow.js レイヤー モデルを使用している場合は、 tfjs-converter を使用してこれをグラフ モデル形式に変換できます。推論の使用例では、グラフ モデル形式の方が効率的です。

tfjs-core を使用した低レベル Tensor 操作

私たちがサポートするもう 1 つのユースケースは、低レベルのテンソル操作に @tensorflow/tjfs-core パッケージを直接使用するプログラムです。

カスタムビルドに対する当社のアプローチ

この機能を設計する際の中心となる原則には次のものが含まれます。

  • JavaScript モジュール システム (ESM) を最大限に利用し、TensorFlow.js のユーザーが同じことをできるようにします。
  • TensorFlow.js を既存のバンドラー (webpack、rollup など) によって可能な限りツリーシェイキング可能にします。これにより、ユーザーはコード分割などの機能を含む、これらのバンドラーのすべての機能を利用できるようになります。
  • バンドルのサイズにそれほど敏感でないユーザーにとっての使いやすさを可能な限り維持します。これは、ライブラリのデフォルトの多くがサイズ最適化されたビルドよりも使いやすさをサポートしているため、実稼働ビルドにはさらに多くの労力が必要になることを意味します。

私たちのワークフローの主な目標は、最適化しようとしているプログラムに必要な機能のみを含む TensorFlow.js 用のカスタムJavaScript モジュールを作成することです。実際の最適化は既存のバンドラーに依存しています。

私たちは主に JavaScript モジュール システムに依存していますが、ユーザー向けコードでモジュール システムを介して指定するのが難しい部分を処理するためのカスタムCLI ツールも提供しています。この 2 つの例は次のとおりです。

  • モデル仕様はmodel.jsonファイルに保存されます
  • 私たちが使用するバックエンド固有のカーネル ディスパッチ システムへの op。

これにより、カスタム 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」が含まれます。
  • モデル: アプリケーションにロードするモデルの 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 モジュールを指すようにできます。バンドル サイズを最大限に節約するには、エイリアスを設定する必要があるモジュールが 3 つあります。

これは、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'),
  }
}

...

ロールアップの同等のコード スニペットは次のとおりです (完全な例はここにあります)。

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'),
    },
  ],
}));

...

バンドラーがモジュールのエイリアシングをサポートしていない場合は、ステップ 3 で作成された生成されたcustom_tfjs.jsから tensorflow.js をインポートするようにimportステートメントを変更する必要があります。Op 定義はツリーシェイクアウトされませんが、カーネルはツリーのままになります。 -震えた。一般に、ツリーシェイキングカーネルは、最終的なバンドルサイズを最大に節約します。

@tensoflow/tfjs-core パッケージのみを使用している場合は、その 1 つのパッケージにエイリアスを付けるだけで済みます。

ステップ 5. バンドルを作成する

バンドラー ( webpackrollup ) を実行してバンドルを生成します。バンドルのサイズは、モジュールのエイリアスなしでバンドラーを実行する場合よりも小さくなるはずです。このようなビジュアライザーを使用して、最終的なバンドルに何が含まれたのかを確認することもできます。

ステップ 6. アプリをテストする

アプリが期待どおりに動作することを必ずテストしてください。