TF.js 実践開発レシピ - Python Keras経験者向け：TensorFlow.jsで画像認識モデルの特定レイヤーをフリーズし効率的にファインチューニングする方法

Python Keras経験者向け：TensorFlow.jsで画像認識モデルの特定レイヤーをフリーズし効率的にファインチューニングする方法

Tags: TensorFlow.js, 画像認識, 転移学習, ファインチューニング, Keras

「TF.js 実践開発レシピ」をご覧いただきありがとうございます。本記事では、TensorFlow.jsを使用して、事前学習済みの画像認識モデルの特定レイヤーをフリーズ（凍結）およびフリーズ解除し、独自のデータセットで効率的にファインチューニングを行う具体的な方法を解説します。PythonでKerasやTensorFlowを使ったモデル開発の経験がある読者の皆様が、WebブラウザやNode.js環境で同様のテクニックを応用する際の参考になれば幸いです。

はじめに：転移学習とレイヤーのフリーズ

画像認識タスクにおいて、ゼロからモデルを学習させるには大量のデータと計算リソースが必要になることが一般的です。そこで有効な手段となるのが転移学習（Transfer Learning）です。これは、ImageNetのような大規模なデータセットで学習済みのモデルを基盤とし、最終層などを独自のタスクに合わせて変更・再学習させる手法です。

さらに、転移学習の応用として、基盤となるモデルの特定レイヤーをフリーズし、学習対象から外すことで、学習効率を高め、 overfitting（過学習）を抑制することがあります。PythonのKerasでは、モデルやレイヤーの trainable プロパティを False に設定することで、そのレイヤーの重みが学習中に更新されないように制御できます。

TensorFlow.jsでも、Python Kerasと同様の概念とAPIを用いて、モデルやレイヤーの trainable プロパティを操作することが可能です。これにより、ブラウザやNode.js環境でも、より柔軟かつ効率的な転移学習やファインチューニングの実装が可能になります。

本記事では、具体的なコード例とともに、以下の内容を順を追って解説します。

TensorFlow.jsで事前学習済みモデルをロードし、レイヤー構成を確認する方法
特定レイヤーをフリーズする方法
フリーズしたモデルに新しい分類層を追加し、ファインチューニングを行う方法
必要に応じてフリーズを解除し、モデル全体または一部をさらに微調整する方法

事前準備：TensorFlow.jsのインストールとモデルのロード

本記事のコード例を実行するには、TensorFlow.jsライブラリが必要です。npmまたはyarnを使ってプロジェクトにインストールしてください。

npm install @tensorflow/tfjs
# または
yarn add @tensorflow/tfjs

ブラウザ環境で使用する場合は、HTMLファイルで<script>タグを使って読み込むこともできます。

次に、事前学習済みの画像認識モデルをロードします。ここでは例として、@tensorflow-models/mobilenet パッケージのMobileNetV2モデルを使用します。

// JavaScript または TypeScript

import * as tf from '@tensorflow/tfjs';
import * as mobilenet from '@tensorflow-models/mobilenet';

async function loadModel() {
  console.log('Loading MobileNetV2 model...');
  // mobilenetモデルをロードします。version 2を指定。
  // モデルは内部でTensorFlow.jsのLayers Modelとしてロードされます。
  const model = await mobilenet.load({ version: 2 });
  console.log('Model loaded successfully.');

  // ロードされたモデルはLayers Modelのインスタンスではないため、
  // 内部のbaseモデル（feature extractor）を取得します。
  // このbaseモデルがレイヤー構成を持っています。
  const baseModel = model.base;

  console.log('Base model layers:');
  // Python Kerasのmodel.summary()に相当する情報の一部を表示できます
  // summary()メソッドはMobileNetモデルオブジェクトには直接ありませんが、
  // Layers Modelのインスタンスであれば利用できます。
  // mobilenet.load()が返すオブジェクトは推論用のラッパーなので、
  // その中のLayersModelを取得する必要があります。
  // MobileNetV2の場合、model.baseがtf.LayersModelです。
  if (baseModel instanceof tf.LayersModel) {
    baseModel.summary();
  } else {
    console.log('Base model is not a LayersModel instance.');
    // レイヤーリストを直接表示する場合
    // baseModel.layers.forEach((layer, index) => {
    //   console.log(`Layer ${index}: ${layer.name} (Trainable: ${layer.trainable})`);
    // });
  }

  return baseModel; // 以降はこのbaseModelに対して操作を行います
}

// モデルをロードして確認
// loadModel(); // 実行する場合はコメントアウトを外す

上記のコードでは、@tensorflow-models/mobilenet を使用してモデルをロードしています。このライブラリが返すオブジェクトは、推論を容易に行うためのラッパーですが、その内部に tf.LayersModel のインスタンス（特徴量抽出器として使われる部分）を持っています。model.base がその tf.LayersModel です。Python Kerasの model.summary() のように、この baseModel に対して summary() メソッドを呼び出すことで、レイヤーの名前、形状、パラメータ数、そして trainable プロパティの状態などを確認できます。

特定レイヤーのフリーズ（Trainableプロパティの設定）

TensorFlow.jsの tf.LayersModel インスタンスが持つ各レイヤーオブジェクトには、trainable というプロパティがあります。これを false に設定することで、そのレイヤーの重みが学習中に更新されないようにできます。

Python Kerasでは model.trainable = False とするとモデル全体のレイヤーがフリーズされますが、TF.jsの LayersModel では、個々のレイヤーの trainable プロパティを制御するのが一般的です。

通常、転移学習ではモデルの浅い層（汎用的な特徴を学習している層）をフリーズし、深い層（タスク固有の特徴を学習している層）を再学習させます。画像認識モデルの場合、畳み込み層（Conv2D）などの特徴抽出を行う層をフリーズし、全結合層（Dense）などの分類を行う層を再学習させることが多いです。

以下のコードは、ロードしたMobileNetV2モデルのベース部分（特徴量抽出器）の全レイヤーをフリーズする例です。

// baseModel は tf.LayersModel のインスタンスを想定

function freezeLayers(model: tf.LayersModel, freeze: boolean = true) {
  model.layers.forEach(layer => {
    layer.trainable = !freeze; // freeze=trueならtrainable=falseに設定
    console.log(`Layer ${layer.name} trainable set to ${layer.trainable}`);
  });
  console.log(`All layers trainable set to ${!freeze}.`);
}

// モデルをロード後、全レイヤーをフリーズする場合
// async function example() {
//   const baseModel = await loadModel(); // 前述のloadModel関数
//   freezeLayers(baseModel, true);
//
//   // フリーズ後のtrainable状態を確認（summaryなどを再度実行）
//   baseModel.summary();
// }
// example(); // 実行する場合はコメントアウトを外す

特定のレイヤーだけをフリーズしたい場合は、レイヤーの名前やインデックスを使って対象を選択します。例えば、最後の畳み込みブロックより前の全てのレイヤーをフリーズし、最後のブロックと続く層は再学習対象とする、といった戦略が考えられます。

MobileNetV2のようなモデルの場合、通常は最後のプーリング層やその前の畳み込み層までをフリーズし、その後に新しい分類器（Dense層など）を追加して学習させます。

// baseModel は tf.LayersModel のインスタンスを想定
// 特定のレイヤーより前をフリーズする例

function freezeLayersUpToIndex(model: tf.LayersModel, maxIndexToFreeze: number) {
  model.layers.forEach((layer, index) => {
    if (index <= maxIndexToFreeze) {
      layer.trainable = false;
      // console.log(`Layer ${layer.name} (Index ${index}) trainable set to false.`);
    } else {
      layer.trainable = true;
      // console.log(`Layer ${layer.name} (Index ${index}) trainable set to true.`);
    }
  });
  console.log(`Layers up to index ${maxIndexToFreeze} frozen. Others are trainable.`);
}

// 例: MobileNetV2の特定のインデックスまでをフリーズ
// 実際のインデックスはmodel.summary()などで確認が必要です
// 例えば、末尾のGlobalAveragePooling2D層（index 155）までをフリーズする場合
// async function exampleFreezeSpecific() {
//   const baseModel = await loadModel();
//   freezeLayersUpToIndex(baseModel, 155); // 例としてのインデックス
//   baseModel.summary(); // trainableの状態が変わっていることを確認
// }
// exampleFreezeSpecific(); // 実行する場合はコメントアウトを外す

どのインデックスまでをフリーズするかは、モデルの構造と転移学習の目的によって異なります。model.summary() を実行してレイヤー構成を確認し、適切なポイントを選択してください。

フリーズしたモデルへの分類層の追加とファインチューニング

ベースモデル（特徴量抽出器）の必要なレイヤーをフリーズしたら、その出力に新しい分類層を追加して、新しいモデルを構築します。これは、Python Kerasの Functional API や Sequential API を使う感覚と似ています。

ここでは、Sequential APIライクにモデルを構築する例を示します。ベースモデルの出力を受け取り、Dropout層や新しいDense層（分類器）を追加する新しい tf.LayersModel を作成します。

// baseModel はフリーズ設定済みの tf.LayersModel インスタンスを想定
// numClasses は新しいタスクのクラス数

function buildFineTuningModel(baseModel: tf.LayersModel, numClasses: number): tf.LayersModel {
  // 新しいSequentialモデルを作成
  const model = tf.sequential();

  // ベースモデル（特徴量抽出器）を追加
  // baseModelのレイヤーのtrainable設定はここで引き継がれます
  model.add(baseModel);

  // 分類器を追加
  // 必要に応じてGlobalAveragePooling2D層を追加することも多いです
  // MobileNetV2のbaseModelは末尾がConv2D層なので、プーリング層が必要
  model.add(tf.layers.globalAveragePooling2D({}));

  // Dropout層を追加（過学習抑制のため）
  model.add(tf.layers.dropout({ rate: 0.5 })); // ドロップアウト率はお好みで調整

  // 新しい分類層（Dense層）を追加
  // ソフトマックス活性化関数を使用してクラス確率を出力
  model.add(tf.layers.dense({ units: numClasses, activation: 'softmax' }));

  console.log('Fine-tuning model architecture:');
  model.summary();

  return model;
}

// モデルの構築とコンパイルの例
// async function buildAndCompileExample() {
//   const baseModel = await loadModel();
//   freezeLayersUpToIndex(baseModel, 155); // 例としてインデックス155までフリーズ

//   const numClasses = 10; // 例：新しいタスクのクラス数
//   const fineTuningModel = buildFineTuningModel(baseModel, numClasses);

//   // モデルのコンパイル
//   fineTuningModel.compile({
//     optimizer: tf.train.adam(0.001), // 学習率はお好みで調整
//     loss: 'categoricalCrossentropy', // 多クラス分類の場合
//     metrics: ['accuracy']
//   });

//   console.log('Model compiled and ready for training.');

//   // ここでデータセットを準備し、model.fit()で学習を開始します
//   // データセットの準備（画像の前処理など）については別記事を参照してください
// }
// buildAndCompileExample(); // 実行する場合はコメントアウトを外す

buildFineTuningModel 関数では、tf.sequential() で新しいモデルを作成し、フリーズ設定済みの baseModel を最初のレイヤーとして追加しています。その後に、新しい分類層を構成するレイヤー（globalAveragePooling2D, dropout, dense）を追加しています。

モデルを学習可能にするためには、compile() メソッドを呼び出す必要があります。ここでオプティマイザ、損失関数、評価指標を指定します。Python Kerasと同様に、Adamなどの様々なオプティマイザや、categoricalCrossentropy といった標準的な損失関数を利用できます。学習率（ここでは0.001）は、特に転移学習やファインチューニングにおいて重要なハイパーパラメータです。一般的に、事前学習済みモデルを基にする場合は、ゼロから学習させる場合よりも低い学習率が推奨されます。

この fineTuningModel に対して、準備したカスタムデータセットを用いて model.fit() メソッドを呼び出すことで、フリーズされていない層（新しく追加した分類層と、もしフリーズ解除していればベースモデルの一部）の重みが更新され、モデルがカスタムタスクに適応されます。

フリーズ解除による微調整（Optional）

新しい分類層で十分に学習が進んだ後、さらにモデルの精度を高めたい場合や、カスタムデータセットが比較的大きい場合は、ベースモデルの一部または全体のフリーズを解除して、さらに低い学習率で微調整（Fine-tuning）を行うことがあります。

これは、ベースモデルのレイヤーが、カスタムデータセットの特定の特徴に合わせてさらに適応できるようにするためです。フリーズ解除を行う場合は、通常、学習率を大きく下げることが重要です。高い学習率のままフリーズ解除すると、事前学習で得られた有用な特徴表現が破壊されてしまう可能性があります。

フリーズを解除するには、再び対象レイヤーの trainable プロパティを true に設定します。

// fineTuningModel は前述のbuildFineTuningModelで構築したモデルを想定
// numClasses はタスクのクラス数

async function unfreezeAndRetrain(fineTuningModel: tf.LayersModel, numClasses: number) {
  console.log('Unfreezing some layers for further fine-tuning...');

  // fineTuningModelの最初のレイヤーがbaseModelです
  const baseModel = fineTuningModel.layers[0] as tf.LayersModel;

  // 例: ベースモデルの最後の数層だけフリーズを解除する
  // どの層を解除するかはモデル構造と実験により決定します
  const startIndexToUnfreeze = 100; // 例としてインデックス100から解除
  baseModel.layers.forEach((layer, index) => {
    if (index >= startIndexToUnfreeze) {
      layer.trainable = true;
      // console.log(`Layer ${layer.name} (Index ${index}) trainable set to true.`);
    }
  });
  console.log(`Layers from index ${startIndexToUnfreeze} in base model unfrozen.`);

  // モデルのコンパイル（学習率を下げて再コンパイルが必要）
  const lowerLearningRate = 0.0001; // 最初の学習率より低い値を設定
  fineTuningModel.compile({
    optimizer: tf.train.adam(lowerLearningRate),
    loss: 'categoricalCrossentropy',
    metrics: ['accuracy']
  });

  console.log('Model re-compiled with lower learning rate.');
  fineTuningModel.summary(); // trainableの状態が変わっていることを確認

  // ここで再びデータセットを用いてmodel.fit()を実行し、微調整を行います
  // 通常、最初の学習よりも少ないエポック数で十分なことが多いです
}

// 最初の学習が完了した後で実行する例
// async function fullFineTuningProcess() {
//   const baseModel = await loadModel();
//   freezeLayersUpToIndex(baseModel, 155); // 最初のフリーズ
//   const numClasses = 10;
//   const fineTuningModel = buildFineTuningModel(baseModel, numClasses);

//   fineTuningModel.compile({
//     optimizer: tf.train.adam(0.001),
//     loss: 'categoricalCrossentropy',
//     metrics: ['accuracy']
//   });

//   console.log('Starting initial training with frozen layers...');
//   // ダミーデータでの学習例 (実際のデータに置き換えてください)
//   const dummyXs = tf.randomNormal([32, 224, 224, 3]); // バッチサイズ32, 画像サイズ224x224x3
//   const dummyYs = tf.zeros([32, numClasses]); // ダミーのラベル
//   await fineTuningModel.fit(dummyXs, dummyYs, {
//     epochs: 5, // 最初の学習のエポック数
//     batchSize: 32,
//     callbacks: { onEpochEnd: (epoch, logs) => console.log(`Epoch ${epoch+1}: loss = ${logs.loss}, accuracy = ${logs.acc}`) }
//   });
//   dummyXs.dispose(); // メモリ解放
//   dummyYs.dispose();

//   console.log('Initial training finished. Starting fine-tuning...');
//   await unfreezeAndRetrain(fineTuningModel, numClasses);

//   // 微調整のための学習
//   // 実際のデータセットと適切なエポック数で学習を実行
//   const dummyXsRetrain = tf.randomNormal([16, 224, 224, 3]); // 小さなバッチサイズでも良い
//   const dummyYsRetrain = tf.zeros([16, numClasses]);
//   await fineTuningModel.fit(dummyXsRetrain, dummyYsRetrain, {
//     epochs: 3, // 微調整のエポック数
//     batchSize: 16,
//     callbacks: { onEpochEnd: (epoch, logs) => console.log(`Retrain Epoch ${epoch+1}: loss = ${logs.loss}, accuracy = ${logs.acc}`) }
//   });
//   dummyXsRetrain.dispose();
//   dummyYsRetrain.dispose();

//   console.log('Fine-tuning finished.');

//   // モデルの保存や推論などに進めます
// }
// fullFineTuningProcess(); // 実行する場合はコメントアウトを外す

フリーズ解除と再コンパイルを行った後は、再度 model.fit() を呼び出して学習を続行します。この際のデータセットは、最初の学習と同じもので構いませんが、通常、学習率を下げているため、より多くのエポック数が必要になる場合もありますが、過学習を防ぐために早期停止（Early Stopping）などのコールバックを利用することも有効です。

実践的な考慮事項

どのレイヤーをフリーズするか: これは実験的に決定することが多いです。一般的には、データセットのサイズが小さいほど、フリーズする層を多くします。データセットが大きい場合は、より深い層までフリーズを解除して微調整することで性能が向上する可能性があります。
学習率: ファインチューニングでは、ゼロから学習させる場合よりも低い学習率を使用することが必須です。特にフリーズ解除して微調整を行う際は、最初の学習率からさらに桁違いに低い値（例: 1/10〜1/100）に設定します。
データセットの準備: モデルに入力する画像データは、事前学習時と同じ前処理（例: サイズ変更、ピクセル値の正規化）を行う必要があります。これはTensorFlow.jsでの画像前処理に関する別の記事で詳しく解説されています。
メモリ管理: ブラウザ環境での学習は、使用可能なメモリ量に制限があります。特に大きなモデルや大きなバッチサイズで学習を行うと、メモリ不足のエラーが発生する可能性があります。tf.dispose() や tf.tidy() を適切に使用して、不要になったテンソルを解放することが重要です。
計算リソースと学習時間: ブラウザ上での学習は、GPUが利用できる環境であれば高速ですが、CPUのみの場合は時間がかかることがあります。Node.js環境でGPUサポートを有効にすることで、より高速な学習が期待できます。

まとめ

本記事では、TensorFlow.jsを使用して事前学習済み画像認識モデルの特定レイヤーをフリーズ・解除し、効率的なファインチューニングを行う方法を解説しました。Python Kerasで trainable プロパティを操作する感覚で、TensorFlow.jsの tf.LayersModel の各レイヤーの trainable プロパティを設定することで、同様の制御が可能であることを示しました。

このテクニックを応用することで、限られたデータセットでも、事前学習済みモデルの強力な特徴抽出能力を活用しつつ、独自のタスクに特化した高性能な画像認識モデルを開発できます。WebブラウザやNode.jsといったJavaScript環境で画像認識AIを実装する際に、ぜひこの手法を活用してみてください。

TensorFlow.jsでは、本記事で紹介した以外にも、様々なモデルアーキテクチャの定義、カスタム損失関数や評価指標の実装、コールバックの利用など、Python Kerasと類似した多くの機能が提供されています。これらの機能を組み合わせることで、さらに複雑で高度なモデル開発が可能になります。

本記事が、皆様のTensorFlow.jsを使った画像認識開発の一助となれば幸いです。