Worker Module + TensorFlow.js

Firefox 114 で Web Worker の ESM (Worker Module) がサポートされました。 この更新によって Web Worker 内でバンドルサイズの最適化がしやすくなりました。

一番嬉しいのはおそらく TensorFlow.js ではないかなと思います。 TensorFlow.js を使う時、今までは以下のように実装していました。

importScripts("https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@4.6.0/dist/tf.min.js");

let model;
(async () => {
  model = await tf.loadGraphModel("model/model.json");
})();

ESM 化する時は、呼び出し側で type: "module" を宣言した後、

const worker = new Worker("worker.js", { type: "module" });

Worker 側で以下のように呼び出しします。 tfjs-backend-webgl を入れないと速度が大幅に落ちます。

import * as tf from "https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-core@4.6.0/+esm";
import "https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-backend-webgl@4.6.0/+esm";
import { loadGraphModel } from "https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-converter@4.6.0/+esm";

let model;
(async () => {
  model = await loadGraphModel("model/model.json");
})();

import が面倒なので以下のようにも書けます。こちらのほうが ESM らしい。 ビルドサイズは 986KB でした。

import {
  browser,
  cast,
  div,
  expandDims,
  scalar,
  tidy,
  setBackend,
} from "https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-core@4.6.0/+esm";
import "https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-backend-webgl@4.6.0/+esm";
import { loadGraphModel } from "https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-converter@4.6.0/+esm";

let model;
(async () => {
  model = await loadGraphModel("model/model.json");
})();

注意点としては以下の書き方ではビルドサイズが小さくなりません。 いまいち理由はわかっていませんが、CPU のビルドが含まれるからかなあ。

import {
  browser,
  cast,
  div,
  expandDims,
  loadGraphModel,
  scalar,
  tidy,
} from "https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@4.6.0/+esm";

let model;
(async () => {
  model = await loadGraphModel("model/model.json");
})();

Wasm 化の例も書いてみます。 setWasmPaths() したディレクトリに wasm ファイルは設置しておく必要があります。

import * as tf from "https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-core@4.6.0/+esm";
import { setWasmPaths } from "https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-backend-wasm@4.6.0/+esm";
import { loadGraphModel } from "https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-converter@4.6.0/+esm";

let model;
(async () => {
  setWasmPaths("/tegaki-abc/");
  await tf.setBackend("wasm");
  model = await loadGraphModel("model/model.json");
})();

バンドルサイズは 660KB + tfjs-backend-wasm-simd.wasm (425KB) で 1.08MB でした。 WebGL を最適化した場合とサイズはほとんど変わりません。

以下の書き方もできます。先程と同様に tfjs-core と tfjs-converter に分けないとファイルサイズは小さくなりません。

import {
  browser,
  cast,
  div,
  expandDims,
  scalar,
  tidy,
  setBackend,
} from "https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-core@4.6.0/+esm";
import { setWasmPaths } from "https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-backend-wasm@4.6.0/+esm";
import { loadGraphModel } from "https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-converter@4.6.0/+esm";

let model;
(async () => {
  setWasmPaths("/tegaki-abc/");
  await tf.setBackend("wasm");
  model = await loadGraphModel("model/model.json");
})();

小さな CNN の MNIST モデルだと推論時間は、simd-wasm = webgl >> wasm でした。 threaded-simd-wasm は面倒で調べてませんが、simd-wasm より早いでしょう。ただし cross-origin isolation の設定とデバッグが面倒な問題があります。 simd-wasm と webgl の平均的な推論速度は大差ありませんでしたが、WebGL は初期ロード時だけはシェーダの設定でやたら遅くなります。 ただ初期ロードを事前に行っておけばデメリットはほとんどなくなります。

結論としては、ESM 化は 400KB ほどファイルサイズが小さくなるのでアリだと思いますが、正直あまり旨味はない気がします。 またよほどのことがない限り Wasm より WebGL のほうが楽なので良いと思いました (2024-04-26)。 Web NN が標準化されたり、WebGPU が安定化したり、Wasm の custom build が進化したり、cross-origin isolation に進展があったかで再検討するのが良さそうです。