TensorFlow.js バックエンド
速度比較
CPU / WebGL / WebGPU

TensorFlow.js は同じコードで CPU・ WebGL・WebGPU の3つのバックエンドを切り替えられます。 しかし「GPU の方が速い」という直感は、演算の種類やデータサイズによっては 正反対の結果をもたらすことがあります。

この記事では Dense・LSTM・Conv2D・Activation など演算の種類を変えながら 3バックエンドの実行時間を計測し、どの演算でどのバックエンドを選ぶべきかを データで示します。

3バックエンドの違いは「どこで・どうやってテンソル演算を実行するか」です。 データの転送コストと並列度のバランスが、速さを決定します。

TensorFlow.js でバックエンドを切り替えながら計測する実装のポイントを解説します。

// バックエンドを切り替えるたびに setBackend → ready を呼ぶ
async function switchBackend(name) {
  // WebGPU は navigator.gpu で事前チェック
  if (name === 'webgpu' && !navigator.gpu) {
    return false;
  }
  try {
    await tf.setBackend(name);
    await tf.ready();
    // 実際に切り替わったか確認（フォールバック検出）
    return tf.getBackend() === name;
  } catch (e) {
    return false;
  }
}

setBackend() だけでは実際に切り替わらない場合があります。 tf.getBackend() で確認するのが確実です。

async function measure(fn, warmupRuns, measureRuns) {
  // ウォームアップ: JIT・シェーダコンパイルを済ませる
  for (let i = 0; i < warmupRuns; i++) {
    const t = fn();
    await t.data();   // GPU 完了を待つ
    t.dispose();
  }

  // 本計測
  const times = [];
  for (let i = 0; i < measureRuns; i++) {
    const t0 = performance.now();
    const result = fn();
    await result.data();  // GPU 同期（これがないと未完了のまま時刻を取る）
    const t1 = performance.now();
    result.dispose();
    times.push(t1 - t0);
  }

  // 最小値（最速スコア）と中央値を返す
  times.sort((a, b) => a - b);
  return {
    min:    times[0],
    median: times[Math.floor(times.length / 2)],
  };
}

ハイライトの await result.data() が計測精度の核心です。 TensorFlow.js の演算は非同期でGPUにオフロードされるため、 data() でデータを取り出すまで実際の完了時刻が確定しません。 このウェイトがないと、全バックエンドがほぼ同じ（= 非常に小さい）時間を返してしまいます。

const OPS = [
  {
    name: 'Dense (MatMul)',
    fn: (sz) => () => {
      const x = tf.randomNormal([sz.batch, sz.in]);
      const w = tf.randomNormal([sz.in, sz.out]);
      return tf.matMul(x, w);
    },
  },
  {
    name: 'LSTM Cell',
    fn: (sz) => () => {
      const model = tf.sequential();
      model.add(tf.layers.lstm({ units: sz.units, inputShape: [sz.seq, sz.feat] }));
      const x = tf.randomNormal([sz.batch, sz.seq, sz.feat]);
      return model.predict(x);
    },
  },
  {
    name: 'Conv2D',
    fn: (sz) => () => {
      const x = tf.randomNormal([sz.batch, sz.h, sz.w, sz.ch]);
      const f = tf.randomNormal([3, 3, sz.ch, sz.filters]);
      return tf.conv2d(x, f, 1, 'same');
    },
  },
  {
    name: 'Softmax',
    fn: (sz) => () => tf.softmax(tf.randomNormal([sz.batch, sz.classes])),
  },
];

各演算を関数として定義しておき、バックエンドを切り替えながら同じ関数を呼ぶ設計にすることで、 公平な比較が可能になります。演算ごとに「サイズパラメータ」を注入できる構造です。

// tf.tidy() でスコープ内のテンソルを自動解放
const result = tf.tidy(() => {
  const x = tf.randomNormal([batch, n]);
  const w = tf.randomNormal([n, n]);
  return tf.matMul(x, w);  // x, w は tidy 終了時に自動 dispose
});

// result だけ手動 dispose（tidy の外に出ているため）
await result.data();
result.dispose();

// テンソル数の確認（ベンチマーク中はこれで監視する）
console.log(tf.memory().numTensors);

ベンチマークは同じ演算を何十回も繰り返すため、テンソルの dispose() 忘れが メモリリークとなり後の計測を汚染します。tf.tidy() で スコープ管理するか、明示的に dispose() するのを徹底してください。

ベンチマーク結果から導かれる、実用的なバックエンド選択の考え方を整理します。

以上をまとめると、実践的な選択指針は以下の通りです。

このベンチマークを iPad mini（A17 Pro チップ）/ Safari + WebGPU 環境で実行した 実測値を掲載します。モバイル最高クラスのSoCでも、演算の種類によって 最速バックエンドが大きく異なることが確認できます。