TF.js 実践開発レシピ

TensorFlow.jsで学習済み画像認識モデルの構造を確認する:Python Kerasのmodel.summary()に相当する機能と実装

Tags: TensorFlow.js, 画像認識, モデル構造, Keras, JavaScript

はじめに

Python Kerasで機械学習モデルを開発している方にとって、model.summary()メソッドはモデルの構造、各レイヤーの出力形状、パラメータ数を一目で把握するための非常に便利な機能です。TensorFlow.jsに移行し、Pythonで学習したモデルをロードしたり、JavaScriptでモデルを定義したりする際に、「このモデルはどのようなレイヤー構成になっているのか」「各レイヤーの出力形状は意図通りか」「モデル全体のパラメータ数はどれくらいか」といった情報を確認したい場面が多々あります。

しかし、TensorFlow.jsのtf.LayersModelには、Python Kerasのmodel.summary()と全く同じ機能を持つメソッドは直接提供されていません。では、TensorFlow.jsで同様の情報を得るためにはどうすれば良いのでしょうか。

この記事では、TensorFlow.jsを用いて学習済み画像認識モデルの構造情報(レイヤー構成、出力形状、パラメータ数など)を取得し、Python Kerasのmodel.summary()に相当する情報を表示する方法について、具体的なコード例を交えて詳細に解説します。Pythonでの開発経験を活かしつつ、TF.jsの世界でのモデル理解を深めたいと考えている読者の方々にとって、実践的な情報を提供することを目指します。

TensorFlow.jsにおけるLayers Modelの構造確認

TensorFlow.jsのtf.LayersModelオブジェクトは、Python Kerasのtf.keras.Modeltf.keras.Sequentialに対応するクラスです。ロードされたtf.LayersModelは、その構成要素である各レイヤーの情報を保持しています。この情報を取得することで、モデルの構造をプログラム的に確認することが可能です。

model.layersプロパティの利用

tf.LayersModelオブジェクトは、layersというプロパティを持っており、これはモデルを構成する全てのレイヤーの配列です。各レイヤーオブジェクトからは、そのタイプ、名前、入出力形状、パラメータ数などの情報を取得できます。

まずは、簡単なLayers Modelを定義するか、既存のモデルをロードする例を考えます。ここでは、TensorFlow.jsに標準で含まれる事前学習済みモデルであるMobileNetV2をロードしてみましょう。

import * as tf from '@tensorflow/tfjs';

async function loadModelAndGetLayers() {
  // MobileNetV2 モデルをロードします
  console.log('Loading MobileNetV2...');
  const model = await tf.loadLayersModel('https://storage.googleapis.com/tfjs-models/tfjs/mobilenet_v2_1.0_224/model.json');
  console.log('Model loaded.');

  // モデルを構成するレイヤーの配列を取得します
  const layers = model.layers;

  console.log(`Total number of layers: ${layers.length}`);

  // 各レイヤーの基本情報を表示します
  layers.forEach((layer, index) => {
    console.log(`Layer ${index}:`);
    console.log(`  Name: ${layer.name}`);
    console.log(`  Class Name: ${layer.className}`); // Conv2D, BatchNorm etc.
    // outputShape は非同期でないと取得できない場合がありますが、ここでは簡単にプロパティとしてアクセス
    console.log(`  Output Shape: ${layer.outputShape}`);
  });

  // モデル全体のパラメータ数を取得します
  console.log(`Total parameters: ${model.countParams()}`);

  // モデルを解放します(メモリ管理のため重要)
  model.dispose();
  console.log('Model disposed.');
}

loadModelAndGetLayers();

上記のコードでは、ロードしたモデルオブジェクトのlayersプロパティからレイヤー配列を取得し、各レイヤーのnameclassNameoutputShapeといったプロパティを表示しています。また、model.countParams()メソッドを使用してモデル全体のパラメータ数を取得しています。

outputShapeプロパティは、レイヤーが構築(build)された後に利用可能になります。事前学習済みモデルのように、ロード時点で入力形状が確定している場合はすぐに利用できます。自身でLayers Modelを定義する場合は、model.build(inputShape)を実行するか、初めてデータを通して推論を行う必要があります。

Keras summary()に相当する情報の取得

Python Kerasのmodel.summary()は、レイヤーのインデックス、タイプ、名前、出力形状、そしてパラメータ数を整形されたテーブル形式で表示します。TensorFlow.jsでこれに相当する情報を取得するには、各レイヤーオブジェクトの以下のプロパティやメソッドを利用します。

これらの情報を用いて、Kerasのsummary()に近い形式で出力する関数を作成してみましょう。

import * as tf from '@tensorflow/tfjs';

async function logModelSummary(model: tf.LayersModel) {
  console.log('Model Summary:');
  console.log('----------------------------------------------------------------');
  console.log('Layer (type)               Output Shape         Param #');
  console.log('================================================================');

  let totalParams = 0;
  let trainableParams = 0;
  let nonTrainableParams = 0;

  for (let i = 0; i < model.layers.length; i++) {
    const layer = model.layers[i];
    const params = await layer.countParams(); // 非同期でパラメータ数を取得

    // レイヤータイプと名前を整形
    const typeName = `${layer.name} (${layer.className})`;
    const paddedTypeName = typeName.padEnd(35); // 適当な長さにパディング

    // 出力形状を整形
    // outputShape は配列なので文字列に変換
    const outputShapeStr = JSON.stringify(layer.outputShape);
    const paddedOutputShape = outputShapeStr.padEnd(20); // 適当な長さにパディング

    // パラメータ数を整形
    const paramCountStr = params.toLocaleString(); // カンマ区切りで整形
    const paddedParamCount = paramCountStr.padEnd(15); // 適当な長さにパディング

    console.log(`${paddedTypeName}${paddedOutputShape}${paddedParamCount}`);

    totalParams += params;

    // trainableParamsとnonTrainableParamsの区別は、layer.trainable プロパティで判断可能
    // ただし countParams は学習可能/不可に関わらずそのレイヤーの全パラメータ数を返します
    // Keras summary()のような詳細な内訳表示のためには、別途 trainable プロパティを考慮する必要があります
    // ここでは簡潔に totalParams のみ表示します
    // より詳細な内訳が必要な場合は layer.weights を確認し、trainable プロパティを持つ weight のパラメータを合計する必要があります
  }

  console.log('================================================================');
  console.log(`Total params: ${totalParams.toLocaleString()}`);
  // ここで trainable/nonTrainable の内訳も表示する場合は、上記ループ内で layer.trainable を確認し集計します
  console.log('----------------------------------------------------------------');
}

async function runSummaryExample() {
  console.log('Loading MobileNetV2 for summary...');
  const model = await tf.loadLayersModel('https://storage.googleapis.com/tfjs-models/tfjs/mobilenet_v2_1.0_224/model.json');
  console.log('Model loaded.');

  await logModelSummary(model); // 作成した summary 関数を呼び出し

  model.dispose();
  console.log('Model disposed.');
}

runSummaryExample();

このlogModelSummary関数は、ロードしたtf.LayersModelオブジェクトを受け取り、各レイヤーの名前、タイプ、出力形状、パラメータ数を取得して、Kerasのsummary()に近い形式でコンソールに出力します。layer.countParams()が非同期メソッドであるため、関数全体もasyncとし、awaitを使用している点に注意が必要です。

出力形状の表現は、TensorFlow.jsのバージョンやモデルのタイプによって若干異なる場合があります。また、パラメータ数の内訳(学習可能/不可)については、より詳細な実装が必要となりますが、基本的な総パラメータ数はcountParams()で取得できます。

Graph Modelの場合

TensorFlow.jsにはtf.LayersModelとは別に、PythonのSavedModelやGraphDefから変換されるtf.GraphModelというモデル形式も存在します。tf.GraphModelは操作のグラフ構造として表現されており、tf.LayersModelのような明確な「レイヤー」の概念を持たない場合があります。そのため、tf.GraphModelに対して上記のようなlayersプロパティを使った構造確認はできません。

tf.GraphModelの構造を確認したい場合は、変換元であるPython側のモデル情報を参照するか、tf.GraphModelオブジェクトの持つ入力/出力テンソル情報などを確認することになります。Graph Modelは主に推論用途で使われることが多く、モデル内部のレイヤー詳細を頻繁に確認する必要性はLayers Modelほど高くないかもしれません。この記事ではLayers Modelに焦点を当てて解説しました。

実践的な開発における考慮事項

まとめ

この記事では、TensorFlow.jsにおいて、特にtf.LayersModelの構造情報を確認する方法について解説しました。Python Kerasのmodel.summary()に相当する機能は直接提供されていませんが、model.layersプロパティと各レイヤーのname, className, outputShape, countParams()といったプロパティやメソッドを組み合わせることで、レイヤー構成、出力形状、パラメータ数といった重要な情報を取得し、整形して表示することが可能です。

紹介したlogModelSummary関数は、ロードした学習済みモデルが期待通りの構造になっているかを確認したり、モデルのデバッグやパフォーマンス評価を行う上で非常に役立ちます。Python Kerasでの開発経験をお持ちの方がTensorFlow.jsに移行する際に、本記事がTF.jsモデルの理解を深める一助となれば幸いです。

今後は、取得したレイヤー情報に基づいて、特定のレイヤーの出力を取得したり、モデルの一部だけを差し替えて転移学習を行ったりと、さらに実践的な開発に進んでいくことができます。