TensorFlowとPythonで特徴抽出器を作成する方法

TensorFlowでは、buffered prefetching（バッファ付きプリフェッチ）を活用することで特徴抽出器（feature extractor）を作成できます。具体的には、trainable=False を設定することがポイントです。

転移学習の基本概念

少なくとも1つの畳み込み層を含むニューラルネットワークは「畳み込みニューラルネットワーク（CNN）」と呼ばれ、CNNを用いることで学習モデルを構築できます。

画像分類における転移学習の直感的な考え方は次のとおりです。大規模かつ汎用的なデータセットで訓練されたモデルは、視覚世界における汎用モデルとして効果的に機能できます。このようなモデルはすでに特徴マップを学習しているため、ユーザーが大規模データセットを使って大規模モデルをゼロから訓練し直す必要はありません。

TensorFlow Hubとは

TensorFlow Hubは、事前訓練済みのTensorFlowモデルを集めたリポジトリです。これらのモデルを利用すれば、TensorFlowで学習モデルを簡単にファインチューニングできます。

本記事では、tf.kerasと組み合わせてTensorFlow Hubのモデルを使用する方法、およびHubから提供される画像分類モデルの使い方を解説します。まず事前訓練済みの分類器モデルに画像を入力して推論させます。この段階では一切の訓練を必要としません。その後、転移学習によってカスタム画像クラスに対応したモデルへとファインチューニングを行います。

なお、以下のコードはGoogle Colaboratory上で実行しています。Google Colabはブラウザ上でPythonコードを実行できる環境であり、面倒な設定が不要で、GPUにも無料でアクセスできます。ColaboratoryはJupyter Notebookを基盤として構築されています。

コード例

AUTOTUNE = tf.data.AUTOTUNE
train_ds = train_ds.cache().prefetch(buffer_size=AUTOTUNE)
print("The dimensions of data")
for image_batch, labels_batch in train_ds:
    print(image_batch.shape)
    print(labels_batch.shape)
    break

コード引用元：https://www.tensorflow.org/tutorials/images/transfer_learning_with_hub

出力結果

The dimensions of data
(32, 224, 224, 3)
(32,)

解説

• TFHubは、最上位の分類層を持たない状態でモデルを配布しています。
• そのため、これらのモデルは転移学習にそのまま利用できます。
• tfhub.devにある互換性のある任意の画像特徴ベクトルモデルを使用可能です。
• trainable=False を指定することで特徴抽出器を作成できます。
• これにより、特徴抽出層内の変数を凍結（freeze）できます。
• 凍結によって、訓練時には新しい分類層のみが更新されるようになります。
• 特徴抽出器は、各画像に対して長さ1280の特徴ベクトルを返します。