TF.js 実践開発レシピ

TensorFlow.jsにおける画像認識モデルの学習プロセス：Python Kerasとの比較から実践まで

Tags: TensorFlow.js, 画像認識, 学習, Keras, JavaScript, Node.js

TensorFlow.jsは、WebブラウザやNode.js環境で機械学習モデルを実行するための強力なライブラリです。特に画像認識の分野では、Pythonで開発したモデルをフロントエンドやサーバーサイドJavaScriptで活用したいというニーズから、TensorFlow.jsが注目されています。

これまでの記事では、主にPythonで学習済みのモデルをTensorFlow.jsで利用する推論処理や、Layers APIを使ったモデル定義について解説してきました。本記事では、一歩進んでTensorFlow.js環境下で、ゼロから、あるいは既存モデルの一部を改変して画像認識モデルを学習させるプロセスに焦点を当てます。Python Kerasでの学習経験を持つ読者の方々がスムーズにTensorFlow.jsでの学習へと移行できるよう、Pythonとの比較を交えながら、具体的なコード例とともに詳細を解説します。

TensorFlow.jsにおける学習環境の準備

TensorFlow.jsでモデルを学習させるためには、実行環境に応じたパッケージのインストールが必要です。

画像データの準備とTensor形式への変換

モデルを学習させるためには、画像データを数値の多次元配列であるTensorに変換する必要があります。画像認識タスクの場合、データは通常、形状 [batch_size, height, width, channels] の4階Tensorとなります。

Python KerasではImageDataGeneratorのような便利なツールがありますが、TensorFlow.jsではJavaScript環境に合わせてデータを準備する必要があります。

例えば、ブラウザで <canvas> や <img> 要素から画像データを取得し、Tensorに変換する基本的なコードは以下のようになります。

async function loadImageAndConvertToTensor(imgElement) {
  const tensor = tf.browser.fromPixels(imgElement); // HTMLImageElementからTensorを作成
  // モデルの入力形状に合わせて前処理（例: リサイズ、正規化）
  const resized = tf.image.resizeBilinear(tensor, [224, 224]); // 例: 224x224にリサイズ
  const normalized = resized.div(255.0); // 例: 0-1の範囲に正規化
  // バッチ次元を追加
  const batched = normalized.expandDims(0); // [height, width, channels] -> [1, height, width, channels]

  tensor.dispose(); // 中間テンソルは適切に解放
  resized.dispose();
  normalized.dispose();

  return batched;
}

Node.js環境では、ファイルシステムから画像を読み込み、画像処理ライブラリ（例: canvas, sharpなど）を使ってピクセルデータを取得し、Tensorに変換します。

const fs = require('fs');
const jpeg = require('jpeg-js');
const tf = require('@tensorflow/tfjs-node'); // Node.js版TF.jsを使用

function loadJpegAndConvertToTensor(imagePath) {
  const jpegData = fs.readFileSync(imagePath);
  const pixels = jpeg.decode(jpegData, true); // RGBA形式でデコード
  const numChannels = 3; // RGBのみを使用
  const numPixels = pixels.width * pixels.height;
  const values = new Float33Array(numPixels * numChannels);

  // RGBAからRGBを抽出
  let valueIndex = 0;
  for (let i = 0; i < numPixels * 4; i += 4) {
    values[valueIndex++] = pixels.data[i];     // R
    values[valueIndex++] = pixels.data[i + 1]; // G
    values[valueIndex++] = pixels.data[i + 2]; // B
  }

  const shape = [pixels.height, pixels.width, numChannels];
  const tensor = tf.tensor3d(values, shape);

  // モデルの入力形状に合わせて前処理（例: 正規化、リサイズ）
  const normalized = tensor.div(255.0); // 例: 0-1に正規化
  tensor.dispose();

  // バッチ次元を追加
  const batched = normalized.expandDims(0); // [height, width, channels] -> [1, height, width, channels]
  normalized.dispose();

  return batched;
}

実際の学習では、大量のデータを効率的に扱うために、画像ファイルパスやURLのリストと対応するラベルを用意し、オンデマンドで画像を読み込み、前処理とバッチ化を行うパイプラインを構築することが一般的です。Pythonのtf.data.Datasetに相当するものはTF.jsには直接ありませんが、JavaScriptのasync/awaitやジェネレータ関数、あるいは外部ライブラリを活用して同様のデータパイプラインを実装することが可能です。

モデルの定義とコンパイル (model.compile)

モデルの構造はtf.LayersModel（tf.Sequentialやtf.Modelクラスを含む）を使用して定義します。これはPython Kerasのtf.keras.Modelやtf.keras.Sequentialに対応します。

const model = tf.sequential({
  layers: [
    tf.layers.conv2d({
      inputShape: [224, 224, 3],
      kernelSize: 5,
      filters: 8,
      activation: 'relu'
    }),
    tf.layers.maxPooling2d({poolSize: 2, strides: 2}),
    tf.layers.conv2d({
      kernelSize: 5,
      filters: 16,
      activation: 'relu'
    }),
    tf.layers.maxPooling2d({poolSize: 2, strides: 2}),
    tf.layers.flatten(),
    tf.layers.dense({units: 10, activation: 'softmax'}) // 例: 10クラス分類
  ]
});

モデルを学習可能にするためには、オプティマイザ、損失関数、評価指標を設定してコンパイルする必要があります。これはPython Kerasのmodel.compile()メソッドと非常に似ています。

// オプティマイザの選択
// 例: Adamオプティマイザ
const optimizer = tf.train.adam();
// 例: SGDオプティマイザ
// const optimizer = tf.train.sgd(0.01); // 学習率を指定

// 損失関数の選択
// 例: 多クラス分類のクロスエントロピー
// ラベルがOne-hotエンコーディングされている場合
const loss = 'categoricalCrossentropy';
// ラベルが整数インデックスの場合 (PythonのSparseCategoricalCrossentropyに相当)
// const loss = 'sparseCategoricalCrossentropy';

// 評価指標の選択
// 例: 精度
const metrics = ['accuracy'];

// モデルのコンパイル
model.compile({
  optimizer: optimizer,
  loss: loss,
  metrics: metrics
});

// コンパイル後のモデル概要確認 (Pythonのmodel.summary()に相当)
model.summary();

tf.trainモジュールには、Adam, SGD, RMSPropなど、Python版TensorFlowでよく使われる多くのオプティマイザが実装されています。損失関数や評価指標も、一般的なものが文字列指定または関数として利用可能です。Python Kerasの経験者であれば、これらの設定は直感的に理解できるでしょう。

モデルの学習実行 (model.fit)

モデルの学習はmodel.fit()メソッドを呼び出すことで実行されます。これはPython Kerasのmodel.fit()と同様に、入力データ、対応するラベル、エポック数、バッチサイズなどを指定します。

async function trainModel(model, trainXs, trainYs, validationXs, validationYs) {
  const history = await model.fit(trainXs, trainYs, {
    epochs: 50,           // 学習エポック数
    batchSize: 32,        // バッチサイズ
    validationData: [validationXs, validationYs], // バリデーションデータ (オプショナル)
    callbacks: {
      // 学習中のコールバック関数を定義 (下記参照)
      onEpochEnd: (epoch, logs) => {
        console.log(`Epoch ${epoch + 1}: loss = ${logs.loss.toFixed(4)}, accuracy = ${logs.acc.toFixed(4)}`);
        if (logs.val_loss) {
          console.log(`  validation loss = ${logs.val_loss.toFixed(4)}, validation accuracy = ${logs.val_acc.toFixed(4)}`);
        }
      }
      // 早期終了などの組み込みコールバックも指定可能
      // callbacks: [new tf.callbacks.EarlyStopping({ monitor: 'val_loss' })]
    }
  });

  console.log("Training finished.");
  return history;
}

// ダミーデータでの実行例
// const trainXs = tf.randomNormal([100, 224, 224, 3]); // 100枚の訓練画像 (ダミー)
// const trainYs = tf.randomUniform([100, 10], 0, 1, 'float32'); // 100枚のラベル (One-hot, ダミー)
// const validationXs = tf.randomNormal([20, 224, 224, 3]); // 20枚のバリデーション画像 (ダミー)
// const validationYs = tf.randomUniform([20, 10], 0, 1, 'float32'); // 20枚のラベル (One-hot, ダミー)

// trainModel(model, trainXs, trainYs, validationXs, validationYs);

// ダミーテンソルは使い終わったら解放
// trainXs.dispose(); trainYs.dispose(); validationXs.dispose(); validationYs.dispose();

model.fitメソッドは非同期処理(async)として実行されるため、awaitキーワードを使用して完了を待つ必要があります。これにより、特にブラウザ環境では学習中のUIブロックを防ぐことができます。

学習中の進捗確認と制御（Callbacks）

Python Kerasと同様に、TensorFlow.jsのmodel.fitメソッドにはcallbacksオプションを指定できます。これにより、学習中のエポック終了時やバッチ終了時などに任意の処理を実行できます。

また、TensorFlow.jsにはいくつかの組み込みCallbackクラスが提供されています。

学習済みモデルの保存とロード

学習済みのモデルは、後で推論に使用したり、学習を再開したりするために保存することができます。TensorFlow.jsでは、環境に応じて様々な保存・ロード方法が提供されています。保存形式は、モデルの構造と重みを含むmodel.jsonファイルと、重みデータを含むバイナリファイル群です。

実践的な考慮事項

TensorFlow.jsで画像認識モデルを学習させる際には、いくつか考慮すべき点があります。

まとめ

本記事では、TensorFlow.jsを使ってブラウザやNode.js環境で画像認識モデルを学習させるための基本的なステップと、Python Kerasとの比較による理解の助けとなる情報を解説しました。データ準備からモデルの定義、コンパイル、学習実行、コールバックによる制御、そしてモデルの保存・ロードまで、一連のプロセスを具体的なコード例とともに示しました。

Python Kerasでの学習経験を持つ読者の方であれば、TensorFlow.jsのLayers APIや学習API（compile, fitなど）がPython版とよく似ていることに気づかれたかと思います。これにより、既存の機械学習知識を活かしつつ、WebやNode.jsの環境でAIモデルの開発・実行に取り組むことが可能です。

TensorFlow.jsでの学習は、特にブラウザ環境では制約もありますが、特定のユースケース（例: エッジデバイスでの追加学習、インタラクティブなデモ、パーソナライズされたモデル学習など）において非常に強力な選択肢となります。本記事が、TensorFlow.jsによる実践的な画像認識AI開発の一助となれば幸いです。