TensorFlow Hub による転移学習

TensorFlow.org で表示 Google Colab で実行 GitHub でソースを表示 ノートブックをダウンロード TF Hub モデルを参照

TensorFlow Hub は、トレーニング済みの TensorFlow モデルのリポジトリです。

このチュートリアルでは、以下の方法を実演します。

  1. TensorFlow Hub からのモデルを tf.keras で利用する。
  2. TensorFlow Hub からの画像分類モデルを使用する。
  3. 独自の画像クラスのモデルを微調整するためにシンプルな転移学習を行う。

セットアップ

import numpy as np
import time

import PIL.Image as Image
import matplotlib.pylab as plt

import tensorflow as tf
import tensorflow_hub as hub

import datetime

%load_ext tensorboard
2024-01-11 22:09:43.978814: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2024-01-11 22:09:43.978858: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2024-01-11 22:09:43.980398: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered

ImageNet の分類器

ImageNet ベンチマークデータセットで事前トレーニングされた分類器モデルを使用するため、初期トレーニングは不要です!

分類器のダウンロード

TensorFlow Hub から事前トレーニング済みの MobileNetV2 モデルを選択し、Keras レイヤーとして hub.KerasLayer でラップします。ここでは、TensorFlow Hub からであれば、以下のドロップダウンに提供されている Example も含み、互換性のあるどの画像分類器モデルでも構いません。

mobilenet_v2 ="https://tfhub.dev/google/tf2-preview/mobilenet_v2/classification/4"
inception_v3 = "https://tfhub.dev/google/imagenet/inception_v3/classification/5"

classifier_model = mobilenet_v2
IMAGE_SHAPE = (224, 224)

classifier = tf.keras.Sequential([
    hub.KerasLayer(classifier_model, input_shape=IMAGE_SHAPE+(3,))
])

1 枚の画像で実行する

モデルを試すために、画像を 1 枚ダウンロードします。

grace_hopper = tf.keras.utils.get_file('image.jpg','https://storage.googleapis.com/download.tensorflow.org/example_images/grace_hopper.jpg')
grace_hopper = Image.open(grace_hopper).resize(IMAGE_SHAPE)
grace_hopper
Downloading data from https://storage.googleapis.com/download.tensorflow.org/example_images/grace_hopper.jpg
61306/61306 [==============================] - 0s 0us/step

png

grace_hopper = np.array(grace_hopper)/255.0
grace_hopper.shape
(224, 224, 3)

バッチの次元を追加し、画像をモデルに入力します。

result = classifier.predict(grace_hopper[np.newaxis, ...])
result.shape
1/1 [==============================] - 2s 2s/step
(1, 1001)

結果は、1001 要素のベクトルのロジットで、画像の各クラスの確率を評価します。

そのため、トップのクラス ID は argmax を使うことでみつけることができます:

predicted_class = tf.math.argmax(result[0], axis=-1)
predicted_class
<tf.Tensor: shape=(), dtype=int64, numpy=653>

推論結果のデコード

predicted_class ID(653 など)を取り、ImageNet データセットラベルをフェッチして予測をデコードします。

labels_path = tf.keras.utils.get_file('ImageNetLabels.txt','https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt')
imagenet_labels = np.array(open(labels_path).read().splitlines())
plt.imshow(grace_hopper)
plt.axis('off')
predicted_class_name = imagenet_labels[predicted_class]
_ = plt.title("Prediction: " + predicted_class_name.title())

png

シンプルな転移学習

ただし、元の ImageNet データセット(事前トレーニング済みモデルがトレーニングされたデータセット)に含まれないクラスを持つ独自のデータセットを使用してカスタム分類器を作成する場合はどうでしょうか。

これは、以下のようにして行います。

  1. TensorFlow Hub から事前トレーニング済みモデルを選択します。
  2. カスタムデータセットのクラスを認識できるよう、最上位(最後)のレイヤーを保持します。

データセット

この例では、TensorFlow flowers データセットを使用します。

import pathlib

data_file = tf.keras.utils.get_file(
  'flower_photos.tgz',
  'https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz',
  cache_dir='.',
   extract=True)

data_root = pathlib.Path(data_file).with_suffix('')
Downloading data from https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz
228813984/228813984 [==============================] - 1s 0us/step

まず、tf.keras.utils.image_dataset_from_directory を使用して、このデータをディスクの画像データを使ったモデルに読み込みます。これにより、tf.data.Dataset が生成されます。

batch_size = 32
img_height = 224
img_width = 224

train_ds = tf.keras.utils.image_dataset_from_directory(
  str(data_root),
  validation_split=0.2,
  subset="training",
  seed=123,
  image_size=(img_height, img_width),
  batch_size=batch_size
)

val_ds = tf.keras.utils.image_dataset_from_directory(
  str(data_root),
  validation_split=0.2,
  subset="validation",
  seed=123,
  image_size=(img_height, img_width),
  batch_size=batch_size
)
Found 3670 files belonging to 5 classes.
Using 2936 files for training.
Found 3670 files belonging to 5 classes.
Using 734 files for validation.

flowers データセットには 5 つのクラスがあります。

class_names = np.array(train_ds.class_names)
print(class_names)
['daisy' 'dandelion' 'roses' 'sunflowers' 'tulips']

次に、画像モデルに使用される TensorFlow Hub の規則では[0, 1] 範囲の浮動小数点数の入力が期待されるため、tf.keras.layers.Rescaling 前処理レイヤーを使用してこれを達成します。

注意: モデルには、tf.keras.layers.Rescaling レイヤーも含めることができます。トレードオフに関する議論について、前処理レイヤーの操作ガイドをご覧ください。

normalization_layer = tf.keras.layers.Rescaling(1./255)
train_ds = train_ds.map(lambda x, y: (normalization_layer(x), y)) # Where x—images, y—labels.
val_ds = val_ds.map(lambda x, y: (normalization_layer(x), y)) # Where x—images, y—labels.

3 番目に、Dataset.prefetch を使って、バッファリングされたプリフェッチで入力パイプラインを終了します。これで、I/O ブロッキングの問題が生じずにディスクからデータを生成することができます。

これらが、データを読み込む際に使用することが推奨される、いくつかの最も重要な tf.data メソッドです。さらに関心がある場合は、これらのメソッド、ディスクへのデータのキャッシュ方法、およびその他の手法について、tf.data API によるパフォーマンスの改善ガイドをご覧ください。

AUTOTUNE = tf.data.AUTOTUNE
train_ds = train_ds.cache().prefetch(buffer_size=AUTOTUNE)
val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)
for image_batch, labels_batch in train_ds:
  print(image_batch.shape)
  print(labels_batch.shape)
  break
(32, 224, 224, 3)
(32,)
2024-01-11 22:09:57.246325: W tensorflow/core/kernels/data/cache_dataset_ops.cc:858] The calling iterator did not fully read the dataset being cached. In order to avoid unexpected truncation of the dataset, the partially cached contents of the dataset  will be discarded. This can happen if you have an input pipeline similar to `dataset.cache().take(k).repeat()`. You should use `dataset.take(k).cache().repeat()` instead.

分類器で画像をバッチ処理する

分類器で画像をバッチ処理していきます。

result_batch = classifier.predict(train_ds)
92/92 [==============================] - 6s 42ms/step
predicted_class_names = imagenet_labels[tf.math.argmax(result_batch, axis=-1)]
predicted_class_names
array(['daisy', 'coral fungus', 'rapeseed', ..., 'daisy', 'daisy',
       'birdhouse'], dtype='<U30')

これらの予測が画像とどれくらい整合しているかを確認します。

plt.figure(figsize=(10,9))
plt.subplots_adjust(hspace=0.5)
for n in range(30):
  plt.subplot(6,5,n+1)
  plt.imshow(image_batch[n])
  plt.title(predicted_class_names[n])
  plt.axis('off')
_ = plt.suptitle("ImageNet predictions")

png

注意: すべての画像は CC-BY のライセンス下にあります。作成者のリストは LICENSE.txt ファイルをご覧ください。

結果は完全とは決して言えませんが、これらはモデルがトレーニングされたクラスではないこと(「daisy」を除く)考慮すれば、合理的です。

ヘッドレスモデルのダウンロード

TensorFlow Hub は最上位の分類層を含まないモデルも配布しています。これらは転移学習に簡単に利用することができます。

TensorFlow Hub から事前トレーニング済みの MobileNetV2 モデルを選択します。ここでは、TensorFlow Hub からであれば、以下のドロップダウンに提供されている Example も含み、互換性のあるどの画像分類器モデルでも構いません。

mobilenet_v2 = "https://tfhub.dev/google/tf2-preview/mobilenet_v2/feature_vector/4"
inception_v3 = "https://tfhub.dev/google/tf2-preview/inception_v3/feature_vector/4"

feature_extractor_model = mobilenet_v2

hub.KerasLayer を使用して、事前トレーニング済みモデルを Keras レイヤーとしてラップし、特徴量抽出器を作成します。trainable=False 引数を使用して変数を凍結し、トレーニングのみが新しい分類器レイヤーを変更できるようにします。

feature_extractor_layer = hub.KerasLayer(
    feature_extractor_model,
    input_shape=(224, 224, 3),
    trainable=False)

特徴量抽出器は、画像ごとに 1280 長のベクトルを返します(この例では、画像バッチサイズは 32 のママになります)。

feature_batch = feature_extractor_layer(image_batch)
print(feature_batch.shape)
(32, 1280)

上位の分類レイヤーを接合する

モデルを完成するために、特徴量抽出器レイヤーをtf.keras.Sequential モデルにラップし、分類用に全結合レイヤーを追加します。

num_classes = len(class_names)

model = tf.keras.Sequential([
  feature_extractor_layer,
  tf.keras.layers.Dense(num_classes)
])

model.summary()
Model: "sequential_1"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 keras_layer_1 (KerasLayer)  (None, 1280)              2257984   
                                                                 
 dense (Dense)               (None, 5)                 6405      
                                                                 
=================================================================
Total params: 2264389 (8.64 MB)
Trainable params: 6405 (25.02 KB)
Non-trainable params: 2257984 (8.61 MB)
_________________________________________________________________
predictions = model(image_batch)
predictions.shape
TensorShape([32, 5])

モデルのトレーニング

Model.compile を使用してトレーニングプロセスを構成し、tf.keras.callbacks.TensorBoard コールバックを追加してログの作成と保存を行います。

model.compile(
  optimizer=tf.keras.optimizers.Adam(),
  loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
  metrics=['acc'])

log_dir = "logs/fit/" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
tensorboard_callback = tf.keras.callbacks.TensorBoard(
    log_dir=log_dir,
    histogram_freq=1) # Enable histogram computation for every epoch.

次に、Model.fit メソッドを使用して、モデルをトレーニングします。

この例を短くするために、10 エポックだけトレーニングします。後で TensorBoard にトレーニングプロセスを可視化できるよう、TensorBoard コールバックでログを作成して保存します。

NUM_EPOCHS = 10

history = model.fit(train_ds,
                    validation_data=val_ds,
                    epochs=NUM_EPOCHS,
                    callbacks=tensorboard_callback)
Epoch 1/10
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1705011010.092570 1004986 device_compiler.h:186] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.
92/92 [==============================] - 11s 84ms/step - loss: 0.7508 - acc: 0.7282 - val_loss: 0.4367 - val_acc: 0.8638
Epoch 2/10
92/92 [==============================] - 6s 63ms/step - loss: 0.3741 - acc: 0.8706 - val_loss: 0.3494 - val_acc: 0.8787
Epoch 3/10
92/92 [==============================] - 6s 63ms/step - loss: 0.2926 - acc: 0.9046 - val_loss: 0.3185 - val_acc: 0.8910
Epoch 4/10
92/92 [==============================] - 6s 63ms/step - loss: 0.2435 - acc: 0.9281 - val_loss: 0.3046 - val_acc: 0.9033
Epoch 5/10
92/92 [==============================] - 6s 63ms/step - loss: 0.2087 - acc: 0.9397 - val_loss: 0.2972 - val_acc: 0.9046
Epoch 6/10
92/92 [==============================] - 6s 63ms/step - loss: 0.1821 - acc: 0.9499 - val_loss: 0.2926 - val_acc: 0.9074
Epoch 7/10
92/92 [==============================] - 6s 63ms/step - loss: 0.1607 - acc: 0.9608 - val_loss: 0.2894 - val_acc: 0.9060
Epoch 8/10
92/92 [==============================] - 6s 63ms/step - loss: 0.1431 - acc: 0.9649 - val_loss: 0.2868 - val_acc: 0.9046
Epoch 9/10
92/92 [==============================] - 6s 63ms/step - loss: 0.1283 - acc: 0.9710 - val_loss: 0.2846 - val_acc: 0.9046
Epoch 10/10
92/92 [==============================] - 6s 62ms/step - loss: 0.1156 - acc: 0.9751 - val_loss: 0.2827 - val_acc: 0.9060

エポックごとに指標がどのように変化しているかを表示し、他のスカラー値を追跡するために、TensorBoard を起動します。

%tensorboard --logdir logs/fit

推論結果の確認

モデルの予測からクラス名の番号付きリストを取得します。

predicted_batch = model.predict(image_batch)
predicted_id = tf.math.argmax(predicted_batch, axis=-1)
predicted_label_batch = class_names[predicted_id]
print(predicted_label_batch)
1/1 [==============================] - 0s 440ms/step
['roses' 'dandelion' 'tulips' 'sunflowers' 'dandelion' 'roses' 'dandelion'
 'roses' 'tulips' 'dandelion' 'tulips' 'tulips' 'sunflowers' 'tulips'
 'dandelion' 'roses' 'daisy' 'tulips' 'dandelion' 'dandelion' 'dandelion'
 'tulips' 'sunflowers' 'roses' 'sunflowers' 'dandelion' 'tulips' 'roses'
 'roses' 'sunflowers' 'tulips' 'sunflowers']

結果をプロットします

plt.figure(figsize=(10,9))
plt.subplots_adjust(hspace=0.5)

for n in range(30):
  plt.subplot(6,5,n+1)
  plt.imshow(image_batch[n])
  plt.title(predicted_label_batch[n].title())
  plt.axis('off')
_ = plt.suptitle("Model predictions")

png

モデルのエクスポート

モデルのトレーニングが完了したので、後で再利用するために、SavedModel としてエクスポートします。

t = time.time()

export_path = "/tmp/saved_models/{}".format(int(t))
model.save(export_path)

export_path
INFO:tensorflow:Assets written to: /tmp/saved_models/1705011071/assets
INFO:tensorflow:Assets written to: /tmp/saved_models/1705011071/assets
'/tmp/saved_models/1705011071'

SavedModel を再読み込みできることと、モデルが同じ結果を出力できることを確認します。

reloaded = tf.keras.models.load_model(export_path)
result_batch = model.predict(image_batch)
reloaded_result_batch = reloaded.predict(image_batch)
1/1 [==============================] - 0s 68ms/step
1/1 [==============================] - 0s 495ms/step
abs(reloaded_result_batch - result_batch).max()
0.0
reloaded_predicted_id = tf.math.argmax(reloaded_result_batch, axis=-1)
reloaded_predicted_label_batch = class_names[reloaded_predicted_id]
print(reloaded_predicted_label_batch)
['roses' 'dandelion' 'tulips' 'sunflowers' 'dandelion' 'roses' 'dandelion'
 'roses' 'tulips' 'dandelion' 'tulips' 'tulips' 'sunflowers' 'tulips'
 'dandelion' 'roses' 'daisy' 'tulips' 'dandelion' 'dandelion' 'dandelion'
 'tulips' 'sunflowers' 'roses' 'sunflowers' 'dandelion' 'tulips' 'roses'
 'roses' 'sunflowers' 'tulips' 'sunflowers']
plt.figure(figsize=(10,9))
plt.subplots_adjust(hspace=0.5)
for n in range(30):
  plt.subplot(6,5,n+1)
  plt.imshow(image_batch[n])
  plt.title(reloaded_predicted_label_batch[n].title())
  plt.axis('off')
_ = plt.suptitle("Model predictions")

png

次のステップ

SavedModel は、読み込んで推論に使用したり、TensorFlow Lite モデル(オンデバイス機械学習)や TensorFlow.js モデル(JavaScript での機械学習)に変換したりできます。

TensorFlow Hub からのトレーニング済みモデルを画像、テキスト、オーディオ、および動画タスクで使用する方法について、その他のチュートリアルをご覧ください。