はじめに
「機械学習をやりたいけど、Pythonの環境構築やサーバーの用意が面倒」と感じたことはないでしょうか。学習済みモデルをちょっと試したいだけなのに、仮想環境を作ってライブラリを揃えて……という手間だけで疲れてしまうことも珍しくありません。
TensorFlow.jsは、この悩みをブラウザだけで解決してくれるライブラリです。Google Brainチームが開発しており、<script>タグ1本、あるいはnpm installだけで、モデルの構築・学習・推論のすべてをJavaScriptだけで完結できます。WebGLを使ったGPUアクセラレーションにも対応しているので、CPUだけでは重い計算もブラウザ上で現実的な速度で動かせます。
とはいえ、説明を読むより実際にテンソルを動かしてみた方が早いと思います。数値の配列を入力すると、その場でテンソル演算の結果が変わるサンプルを用意したので、先に触ってみたい方はこちらからどうぞ。
TensorFlow.jsとは
TensorFlow.jsは、Googleが開発しているJavaScript/TypeScript向けの機械学習ライブラリです。もともとはdeeplearn.jsという名前のプロジェクトでしたが、TensorFlowファミリーに統合されて現在の形になりました。GitHubで19,000以上のスターを集めており、現在も活発に開発が続いています。
ブラウザ上で新しくモデルを学習させることもできますし、Pythonで訓練済みのTensorFlow/Kerasモデルを変換してブラウザで推論だけ行うこともできます。さらにNode.js上で動かすこともでき、フロントエンドからサーバーサイドまで同じAPIで機械学習を扱えるのが強みです。
主な特徴
- ブラウザで完結する - モデルの学習も推論もクライアントサイドのJavaScriptだけで行える。サーバーへのデータ送信が不要なため、プライバシー面でも有利
- 複数の演算バックエンド - WebGL・WebGPU・WASM・CPUといった複数のバックエンドを切り替えられ、実行環境に応じて最適な速度を選べる
- Pythonモデルとの互換性 - Kerasやtf.SavedModel形式のモデルを
tensorflowjs_converterでTensorFlow.js形式に変換し、そのままブラウザで読み込める - Node.js対応 -
@tensorflow/tfjs-nodeを使えばサーバーサイドでネイティブのTensorFlow C++バインディングによる高速実行も可能
インストール
npm install @tensorflow/tfjs
yarn add @tensorflow/tfjs
pnpm add @tensorflow/tfjs
CDN経由で<script>タグから読み込むこともできます。
<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@4.22.0/dist/tf.min.js"></script>
TensorFlow.jsのサンプルを動かす
TensorFlow.jsの基本はtf.tensor()系のAPIでテンソル（多次元配列）を作り、add()やmul()といったメソッドで演算することです。以下のサンプルでは、2つの配列をtf.tensor1d()でテンソル化し、加算・要素積・平均値を計算しています。
要点だけを抜き出すと、次のような形になります。
import * as tf from '@tensorflow/tfjs'
const a = tf.tensor1d([1, 2, 3, 4])
const b = tf.tensor1d([10, 20, 30, 40])
const sum = a.add(b) // [11, 22, 33, 44]
const product = a.mul(b) // [10, 40, 90, 160]
console.log(await sum.array())
tf.dispose([a, b, sum, product])
実際に動かせるものが下です。入力欄の数値をカンマ区切りで書き換えると、その場でテンソル演算の結果が変わります。
配列Aを1,2,3,4,5のように要素数を増やして書き換えても、配列Bと要素数さえ揃っていればadd()とmul()はそのまま動きます。テンソルは形状（shape）が合っていないと演算エラーになる点も、実際に配列の長さを変えて試すと体感できます。
基本的な使い方
TensorFlow.jsの核となるワークフローは、「モデルを定義する」「compileで学習方法を決める」「fitで学習する」「predictで推論する」の4ステップです。以下はy = 2x - 1という単純な関数を学習させる、最小構成の例です。
import * as tf from '@tensorflow/tfjs'
// y = 2x - 1 を学習するだけの最小構成モデル
const model = tf.sequential()
model.add(tf.layers.dense({ units: 1, inputShape: [1] }))
model.compile({ optimizer: 'sgd', loss: 'meanSquaredError' })
const xs = tf.tensor2d([-1, 0, 1, 2, 3, 4], [6, 1])
const ys = tf.tensor2d([-3, -1, 1, 3, 5, 7], [6, 1])
await model.fit(xs, ys, { epochs: 200 })
const prediction = model.predict(tf.tensor2d([10], [1, 1]))
console.log(await prediction.array()) // 19に近い値
tf.sequential()でレイヤーを積み上げていく形式のモデルを作り、tf.layers.dense()で全結合層を追加します。model.fit()は非同期関数で、指定したエポック数だけ学習を繰り返します。学習が終わったmodel.predict()に新しい入力を渡すと、学習した関係に基づいた予測値が返ってきます。
実践的なユースケース
ニューラルネットワークの学習と予測をリアルタイムに確認する
「学習」と「推論」を別々のコード片として眺めるだけでは、モデルが少しずつ賢くなっていく様子が分かりません。以下のサンプルは、ボタンを押すとmodel.fit()が実行され、callbacks.onEpochEndで各エポックの損失（loss）が画面に反映されます。学習後に任意のxの値を入力すると、model.predict()の結果が正解（2x - 1）にどれだけ近いかを確認できます。
「学習を開始」を押すとlossの値が徐々に0へ近づいていくのが分かります。学習後に「予測」を押すと、xにどんな値を入れても2x - 1にかなり近い値が返ってくるはずです。エポック数を減らすとどうなるか、コードを書き換えて試してみるのも面白いところです。
演算バックエンドを切り替えて速度を比較する
TensorFlow.jsはtf.setBackend()で演算の実行環境（バックエンド）を切り替えられます。同じ計算でも、GPUを使うwebglとCPUで計算するcpuとでは速度が大きく変わることがあります。以下のサンプルでは、500×500の行列積（tf.matMul()）を各バックエンドで実行し、かかった時間を比較できます。
webglバックエンドはGPUに計算を送るオーバーヘッドがあるぶん、小さな計算ではcpuの方が速いこともあります。行列のサイズをもっと大きくして試すと、webglが優位になる境目が見えてくるはずです。実行環境（ブラウザやGPU）によっても結果は変わるので、手元のPCとスマホで比べてみるのもおすすめです。
学習したモデルをブラウザに保存・読み込みする
せっかく学習したモデルも、ページをリロードするたびに毎回学習し直すのでは実用的ではありません。TensorFlow.jsはmodel.save()とtf.loadLayersModel()にlocalstorage://というスキームを渡すことで、ブラウザのlocalStorageにモデルの重みとアーキテクチャをそのまま保存・復元できます（本番のWebアプリではindexeddb://を使う方が容量面で安全です）。
「学習して保存」を押した後、ブラウザの開発者ツールでlocalStorageの中身を見ると、tensorflowjs_models/tfjs-demo-modelというキーでモデルの情報が保存されているのが分かります。「読み込んで予測」は学習をやり直さずに、保存済みの重みだけを読み込んで推論しています。実際のWebアプリでオフライン推論やページ再読み込み後の即時推論を実現したいときに使えるパターンです。
まとめ
TensorFlow.jsを使えば、Pythonの環境構築もサーバーの用意もなしに、ブラウザだけでテンソル演算からニューラルネットワークの学習・推論までを完結できます。tf.sequential()とtf.layers.dense()でモデルを組み立て、fit()で学習し、predict()で推論するという流れはPython版のKerasとほぼ同じ感覚で書けるので、機械学習の学習コストを大きく下げてくれます。
WebGL・WASM・CPUといったバックエンドを切り替えられる柔軟性や、localStorage/IndexedDBへのモデル保存など、ブラウザという実行環境ならではの機能も揃っています。まずは今回のサンプルをコピーして、自分の手元のデータで小さなモデルを学習させるところから試してみてください。