事前学習済みモデルを学習させてみた - 茨城エンジニアのPython開発日記

ブログから記事を見つけたい場合はこちら

ブログ地図 - 茨城エンジニアのPython開発日記

こんにちは松原です。ずっと眠いです。

さて、今日は事前学習済みモデルを学習させてみました。
学習済みの特徴量抽出部分を使い、最後の識別部分だけ自分で付け加えて学習させる手法です。

googleが学習させたモデルを利用できるので、性能がでやすいはず。

下記のチュートリアルを参考にしました。

TensorFlow Hub

tensorflow-hubというライブラリを利用して簡単にできます。

チュートリアルをかみ砕いて、必要な部分だけを取り出したり諸々改造したのが下のコード。

#参考URL
#https://www.tensorflow.org/hub?hl=ja
#https://tfhub.dev/s?module-type=image-feature-vector&tf-version=tf2
import tensorflow as tf
import tensorflow_hub as hub

#設定
BATCH_SIZE = 32
EPOCH_NUM = 8
do_fine_tuning = False
IMAGE_SIZE = (224, 224)

#事前学習済みモデル取得
handle_base = "mobilenet_v2_100_224"
MODULE_HANDLE ="https://tfhub.dev/google/imagenet/{}/feature_vector/4".format(handle_base)

#学習データ取得
data_dir = "train_data"

#画像前処理情報
datagen_kwargs = dict(rescale=1./255, validation_split=.20)
dataflow_kwargs = dict(target_size=IMAGE_SIZE, batch_size=BATCH_SIZE,
                   interpolation="bilinear")

#テストデータ用意
test_datagen = tf.keras.preprocessing.image.ImageDataGenerator(
    **datagen_kwargs)
test_generator = test_datagen.flow_from_directory(
    data_dir, subset="validation", shuffle=False, **dataflow_kwargs)

#訓練データ用意
train_datagen = test_datagen
train_generator = train_datagen.flow_from_directory(
    data_dir, subset="training", shuffle=True, **dataflow_kwargs)


#識別器部分追加
model = tf.keras.Sequential([
    tf.keras.layers.InputLayer(input_shape=IMAGE_SIZE + (3,)),
    hub.KerasLayer(MODULE_HANDLE, trainable=do_fine_tuning),
    tf.keras.layers.Dropout(rate=0.2),
    tf.keras.layers.Dense(train_generator.num_classes,
                          kernel_regularizer=tf.keras.regularizers.l2(0.0001))
])
model.build((None,)+IMAGE_SIZE+(3,))

#モデルコンパイル
model.compile(
  optimizer=tf.keras.optimizers.SGD(lr=0.005, momentum=0.9), 
  loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True, label_smoothing=0.1),
  metrics=['accuracy'])

#学習開始準備
steps_per_epoch = train_generator.samples // train_generator.batch_size
validation_steps = test_generator.samples // test_generator.batch_size

#学習実行
model.fit(
    train_generator,
    epochs=EPOCH_NUM, steps_per_epoch=steps_per_epoch,
    validation_data=test_generator,
    validation_steps=validation_steps).history

#モデル保存
tf.saved_model.save(model, "model_save")