TF.js 実践開発レシピ - TensorFlow.jsで画像認識モデルの特定レイヤー出力を取得する方法：Keras/TensorFlow経験者向け解説

TensorFlow.jsで画像認識モデルの特定レイヤー出力を取得する方法：Keras/TensorFlow経験者向け解説

Tags: TensorFlow.js, 画像認識, 中間層出力, モデル解説, Keras

画像認識モデルを扱う際、最終的な推論結果だけでなく、モデルの途中にある特定レイヤーの出力（アクティベーション）を確認したい場合があります。これは、モデルが入力からどのような特徴を抽出しているかを理解したり、デバッグを行ったり、あるいはその中間特徴量を別のタスクに利用したり（例: 特徴量抽出器として使用）する際に非常に有用です。

PythonのTensorFlowやKerasでは、Model APIを使って入力から複数または特定レイヤーの出力を取得する新しいモデルを簡単に構築できます。TensorFlow.jsでも、これと同様の操作を行うための機能が提供されています。本記事では、TensorFlow.jsを用いて学習済みの画像認識モデルから特定レイヤーの出力を取得する具体的な方法について、Pythonでの経験を持つ読者の方々がスムーズに理解できるよう解説します。

なぜ特定レイヤーの出力を取得する必要があるのか

通常のモデルの推論処理（model.predict()など）は、モデルの最終的な出力（例えば分類タスクにおけるクラス確率）を返すように設計されています。しかし、以下のようなケースでは、中間層の出力が必要になります。

モデルの挙動理解: 各レイヤーが入力に対してどのような変換を行い、どのような特徴を捉えているのかを可視化することで、モデルがどのように機能しているのかを深く理解できます。
デバッグ: モデルが期待通りの性能を発揮しない場合、中間層の出力を確認することで、問題が発生している可能性のあるレイヤーを特定する手がかりが得られます。
特徴量抽出: 事前学習済みのモデルを、より高次の特徴量を抽出するための「特徴量抽出器」として使用する際に、特定の畳み込み層の出力などを利用します。転移学習の際にも、最終層を置き換える前に中間層の出力を固定して利用することがあります。
新しいモデルの構築: 特定レイヤーの出力を別のネットワークへの入力とするなど、既存モデルをモジュールとして活用する新しいネットワークを構築する際に必要となります。

Python/Kerasでこれらの操作を行った経験がある方は多いと思いますが、TensorFlow.jsでも同様のことが可能です。

TensorFlow.jsで特定レイヤー出力を取得する基本

TensorFlow.jsにおいて、ロードしたモデルから特定レイヤーの出力を取得する最も一般的な方法は、Python/Kerasと同様に、既存モデルの入力と、取得したい特定レイヤーの出力を結びつけた新しいモデルを構築することです。

この操作は、tf.model() APIを使用して行います。tf.model() は、inputs と outputs を指定して新しいモデルを生成できます。ここでは、既存モデルの入力テンソルを新しいモデルの入力とし、既存モデルの特定レイヤーの出力テンソルを新しいモデルの出力とします。

具体的なコード例：MobileNetV2の中間層出力を取得する

ここでは、TensorFlow.jsで利用可能な事前学習済みモデルである MobileNetV2 を例に、特定レイヤーの出力を取得するコードを示します。

まず、必要なライブラリをインポートし、MobileNetV2 モデルをロードします。

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

async function loadModelAndGetIntermediateOutput(imageElement) {
    console.log('Loading MobileNetV2 model...');
    // MobileNetV2のロード。ここではバージョン2を指定。
    // classifierOnly: trueにすると、最終分類層を除いた特徴量抽出器のみがロードされる
    // load()はデフォルトでバージョン1をロードするので、load(2)とするか、
    // mobilenet.load({version: 2})のようにオプションを指定することが推奨される
    const baseModel = await mobilenet.load({version: 2, alpha: 1.0});
    console.log('MobileNetV2 model loaded.');

    // モデルの構造を確認するためにサマリーを表示することも有効です (ブラウザの開発者コンソールに出力されます)
    // baseModel.model.summary(); // mobilenet.loadで取得したモデルは内部に.modelプロパティを持つ場合がある

    // Mobilenet v2モデル自体のsummaryを確認する場合は、
    // tf.loadGraphModel または tf.loadLayersModel でロードしたモデルに対して summary() を呼び出します。
    // mobilenet.load() は少しラッパーが適用されているため、直接 tf.LayersModel とは異なる場合があります。
    // 特定レイヤーにアクセスするには、baseModel.model.layers を使用します。
    const originalModel = baseModel.model; // ロードしたMobilenetモデルから内部のtf.LayersModelを取得

    // モデルのレイヤー名を確認します。開発者コンソールで originalModel.layers を確認すると良いでしょう。
    // MobileNetV2の場合、中間層には 'block_...' や 'conv_...' といった名前がついています。
    // 例として、'block_13_expand' というレイヤーの出力を取得してみます。
    const targetLayerName = 'block_13_expand';
    const targetLayer = originalModel.getLayer(targetLayerName);

    if (!targetLayer) {
        console.error(`Layer '${targetLayerName}' not found in the model.`);
        return null;
    }

    // 既存モデルの入力と、指定したレイヤーの出力を結びつけた新しいモデルを構築します。
    console.log(`Creating a new model to extract output from layer: ${targetLayerName}`);
    const intermediateModel = tf.model({
        inputs: originalModel.inputs,      // 元モデルの入力を使用
        outputs: targetLayer.output        // 指定したレイヤーの出力を使用
    });
    console.log('Intermediate model created.');

    // 入力画像を準備します。
    // mobilenet.load() でロードしたモデルは、入力として ImageData, HTMLImageElement, HTMLCanvasElement, HTMLVideoElement
    // あるいは tf.Tensor を受け付ける classify() メソッドなどを持っています。
    // ここでは tf.Tensor を直接作成する例を示します。
    // MobilenetV2の入力サイズは 224x224 です。画像をリサイズしてテンソル化します。
    const imageTensor = tf.browser.fromPixels(imageElement)
        .resizeBilinear([224, 224]) // Mobilenetの標準入力サイズにリサイズ
        .toFloat()                 // float型に変換
        .expandDims(0);            // バッチ次元を追加 (モデルは通常 [batch_size, height, width, channels] 形式を期待します)

    // 中間層モデルを使って推論を実行し、出力を取得します。
    console.log('Performing inference with the intermediate model...');
    const intermediateOutput = intermediateModel.predict(imageTensor);
    console.log('Inference complete.');

    // 出力テンソルの形状などを確認します。
    console.log(`Shape of the intermediate output: ${intermediateOutput.shape}`);
    // 出力テンソルの内容をログに出力することも可能ですが、サイズが大きい場合は注意が必要です。
    // console.log('Intermediate output tensor data (first 10 elements):', await intermediateOutput.data());

    // テンソルはGPUメモリなどを消費するため、不要になったら必ず dispose() で解放します。
    imageTensor.dispose();
    intermediateOutput.dispose();
    intermediateModel.dispose(); // 新しく作ったモデルも解放

    console.log('Intermediate output acquired and memory released.');

    return intermediateOutput.shape; // 例として形状を返す
}

// HTML上に <img id="testImage" src="path/to/your/image.jpg" /> があると仮定
// 実際の使用時は、適切な画像要素やcanvas要素を引数として渡してください。
// const imgElement = document.getElementById('testImage');
// if(imgElement) {
//     loadModelAndGetIntermediateOutput(imgElement).then(shape => {
//         console.log('Finished getting intermediate output. Output shape was:', shape);
//     }).catch(error => {
//         console.error('An error occurred:', error);
//     });
// } else {
//    console.error('Image element not found.');
// }

上記のコードでは、以下のステップを実行しています。

mobilenet.load() を使用して MobileNetV2 モデルをロードします。mobilenet.load() は、TensorFlow.jsの低レベルAPIである tf.loadLayersModel や tf.loadGraphModel のラッパーとして機能することがあります。baseModel.model として、基となる tf.LayersModel インスタンスにアクセスできる場合が多いです。
originalModel.getLayer(targetLayerName) を使って、名前で目的のレイヤーを取得します。正確なレイヤー名は、モデルの構造（originalModel.summary() の出力などを参照）を確認して指定する必要があります。
tf.model({ inputs: ..., outputs: ... }) を呼び出し、元のモデルの入力 (originalModel.inputs) と、指定したレイヤーの出力 (targetLayer.output) を持つ新しい tf.LayersModel インスタンス (intermediateModel) を作成します。これが特定レイヤーの出力を得るための新しい実行グラフとなります。
入力画像をテンソルに変換し、MobileNetV2 が期待する形状とデータ型に前処理します。tf.browser.fromPixels、resizeBilinear、toFloat、expandDims といったメソッドを使用します。
作成した intermediateModel の predict() メソッドに前処理済み入力テンソルを渡し、中間層の出力を取得します。
取得した出力テンソルの形状や内容を確認します。
重要: 使用しなくなったテンソル（imageTensor, intermediateOutput）やモデル (intermediateModel) は、dispose() メソッドを呼び出して明示的にメモリを解放する必要があります。JavaScriptのガベージコレクションだけでは、GPUメモリなどが適切に解放されないことがあるため、TF.jsでは手動での dispose() が推奨されます。

Python (Keras/TensorFlow) との比較

PythonでKerasを使って同様の操作を行う場合、以下のようなコードになります。

import tensorflow as tf
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.models import Model

# モデルのロード (学習済み重み込み)
base_model = MobileNetV2(weights='imagenet')

# モデルのサマリーを確認 (レイヤー名を確認)
# base_model.summary()

# 特定のレイヤー名を取得
target_layer_name = 'block_13_expand' # 例

# 指定したレイヤーを取得
target_layer = base_model.get_layer(target_layer_name)

# 新しいモデルを構築 (元の入力と指定レイヤーの出力を結びつける)
intermediate_model = Model(inputs=base_model.inputs, outputs=target_layer.output)

# ダミー入力データの準備 (例: 1枚の 224x224 RGB画像)
dummy_input = tf.random.normal((1, 224, 224, 3))

# 推論の実行
intermediate_output = intermediate_model.predict(dummy_input)

print(f"Shape of the intermediate output: {intermediate_output.shape}")

# モデルの解放 (Pythonでは通常不要ですが、ここでは比較として記載)
# del base_model
# del intermediate_model

ご覧のように、TensorFlow.jsとPython/KerasではAPIの名前（getLayer vs get_layer、model vs Model）やテンソルの生成・操作方法（tf.browser.fromPixels, resizeBilinear など vs tf.random.normal、画像読み込みライブラリの使用など）に違いはありますが、既存モデルの入力と特定レイヤーの出力を結びつけて新しいモデルを構築するという基本的な考え方と手順は非常に類似しています。

Pythonでの経験があれば、TF.jsで特定レイヤー出力を取得する仕組みも比較的容易に理解できるでしょう。

実践的な考慮事項

レイヤー名の確認: モデルの正確なレイヤー名を知るには、model.summary() の出力を確認するのが最も確実です。GraphModelとしてロードしたモデルの場合、summary()メソッドがないことがあるため、モデル構造を別途確認する必要がある場合があります。tf.loadLayersModelでロードした場合はsummary()が利用可能です。
複数レイヤーの出力: 複数のレイヤーの出力を同時に取得したい場合は、outputs をテンソルの配列として指定します。例えば、outputs: [layer1.output, layer2.output] のように記述します。
パフォーマンス: 特定レイヤーの出力を取得するために新しいモデルを作成し実行することは、オリジナルのモデル全体を実行するのと同等か、それ以上の計算コストがかかる可能性があります。特にブラウザ環境では、パフォーマンスへの影響を考慮する必要があります。
メモリ管理: 前述の通り、TensorFlow.jsではテンソルのメモリ管理が重要です。不要になったテンソルやモデルは必ず dispose() してください。特にループ内で頻繁に推論を行う場合は、メモリリークに十分注意が必要です。
WebWorkerの活用: UIスレッドでの計算はブラウザの応答性を損なう可能性があります。計算負荷の高い中間層出力の取得処理は、WebWorker内で実行することを検討してください。

まとめ

本記事では、TensorFlow.jsを用いて学習済みの画像認識モデルから特定レイヤーの出力を取得する方法を解説しました。Python/Kerasでの経験がある方にとっては馴染みやすい、既存モデルの入力と特定レイヤーの出力を結びつけた新しいモデルを構築するというアプローチが基本となります。

中間層の出力を取得する技術は、モデルのデバッグや理解を深めるだけでなく、事前学習済みモデルを強力な特徴量抽出器として活用するなど、様々な応用に繋がります。ぜひ、TensorFlow.jsでの画像認識開発において、このテクニックを活用してみてください。

今後、この記事で取得した中間層出力を実際に可視化する方法や、抽出した特徴量を別のタスクに利用する方法などについても、別の記事で詳しく解説していく予定です。