[python][tensorflow] Tensorflow-hubで学習済みモデルを使ってfine-tuningする

 tensorflow-hubを使うことで、大量のデータを使って学習したモデルを容易に使うことができます。

 今回は例として、学習済みのResnet50から画像の特徴量を抽出して、オリジナルの層を加えて分類を行う方法について書きます。

 ライブラリの説明・インストールなどはこちらを参照してください➡︎https://qiita.com/FukuharaYohei/items/a6fe3c0b0f6e236c4876

実装


import tensorflow_hub as hub
import tensorflow as tf


# tensorflow-hub取り込み。fine-tuningできるよう引数指定
module = hub.Module("https://tfhub.dev/google/imagenet/resnet_v2_50/feature_vector/1", tags={"train"}, trainable=True)
# TF-hubの入力データのサイズ取得
height, width = hub.get_expected_image_size(module)

x = tf.placeholder(tf.float32, (None,) + input_dim)
# TF-hubの入力サイズに対応させる
resized_x = tf.map_fn(lambda x: tf.image.resize_images(x, size=(width, height,)), x)
# 学習済みモデルから特徴量を取得
features = module(resized_x, as_dict=True)['default']
# 分類
prediction_result = \
    predict_layer(tf.tanh(features), output_num, hidden_units, is_training)
t = tf.placeholder(tf.float32, [None, output_num])
cross_entropy = -tf.reduce_mean(t * tf.log(tf.clip_by_value(prediction_result, 1e-12, 1.0)))
train_step = tf.train.AdamOptimizer(learning_rate).minimize(cross_entropy)


def predict_layer(x, output_num, hidden_units, is_training):
    # 推論レイヤ
    keep_prob = tf.cond(is_training, lambda: 0.5, lambda: 1.0)
    output_x = tf.layers.dense(x, units=hidden_units, activation=tf.nn.relu,
                               kernel_initializer=tf.keras.initializers.he_normal())
    output_x = tf.layers.dropout(output_x, keep_prob)
    return tf.layers.dense(output_x, units=output_num, activation=tf.nn.softmax,
                           kernel_initializer=tf.keras.initializers.he_normal())

コードについて解説


 6行目にtags={“train”}、trainable=Trueを指定することで学習済みモデルも学習させることができます。

 URLは任意のTF-hubのFeature Vectorに置き換えできます。https://tfhub.dev/s?module-type=image-feature-vector

 8行目でTF-hubの学習済みモデルが期待する入力データのサイズを取得し、12行目で画像をリサイズしています。tf.map_fnはミニバッチ全体にリサイズを適用するための関数です。

 14行目で学習済みモデルから特徴量を取得します。

 predict_layer関数でオリジナルの層を追加して、得られた特徴量から指定したクラス数に分類します。

コメントを残す

メールアドレスが公開されることはありません。