Technical Hedgehog

機械学習と気になる技術を語るはりねずみ

Googleの事前学習済みモデルを手軽に利用出来るTensorFlow Hub

TensorFlow 機械学習

自然言語処理におけるword2vecや画像処理におけるInceptionなど、一般的に広く用いられているモデルを上流で用いる事は多くあります。汎用的な知識を扱えるメリットがある一方、学習には大量のデータセットの準備と膨大な学習時間がかかってしまいます。

f:id:kamujun:20180810181316p:plain

この問題に対して、あらかじめ学習させた状態のモデル(事前学習済みモデル)を用意しておき上流に転移させる方法があります。本記事ではその事前学習済みモデルについて、Googleが提供するのライブラリであるTensorFlow Hubを紹介します。

TensorFlow HubはGoogleの大量リソースを用いて学習したモデルを手軽に実装できるほか、自作したモデルを別環境で利用しやすいように自作することも可能です。本記事では概要と特徴、利用方法を紹介します。

今回説明するTensorFlow Hubの利用方法、作成方法について実験したコードはGitHubにあります。利用方法は生TensorFlowとKerasにおいてFine-tuningを行う実験を行っています。作成を行う例はネット上でもなかなか無かったので、誰かのお役に立てたら幸いです。(チュートリアルでは大雑把にしか書いてなかった)スターしていただけると嬉しいです。

TensorFlow Hubとは

f:id:kamujun:20180810181604p:plain 画像:google Developers Blog より

TensorFlow HubはGoogleが提供している事前学習済みモデルです。以前から事前学習済みモデルは様々なフレームワークでも提供されており、TensorFlowにおいてもグラフと学習済みチェックポイントが公開されています(現在も更新されているようです)。ただ、利用する際にはグラフを定義したファイルの読み込みやチェックポイントの解凍やロードが必要であるなど手軽ではない印象を受けます。

TensorFlow Hubの特徴はなんといっても手軽に実装できる点が特徴です(しかもGoogleが大量リソースを用いて学習したモデル)。事前学習済させるモデルの多くは大量なコーパスをもとにしているので自作すると大変ですが、TensorFlow HubはGoogleの潤沢なリソースを用いて学習させたモデルが無料で提供されています。例としては画像系の事前学習済みモデルであるNASNetは62,000時間以上のGPU処理時間を投じて計算されたモデルであったり、テキスト系のNeural-Net Language Modelsでは2,000億単語を用いて学習したモデルなどがあります。これらのモデルは後述するモジュールと言う形で数行で実装することが可能なくらい簡単に利用出来ます。

概要・用語説明

TensorFlow Hubにおいてグラフと重みはひとまとめに提供されており、モジュールと呼ばれます。モジュールという単位にすることによってグラフの読み込みや重みのロードなどそれぞれの操作を意識することなく利用することができます。現時点では提供されているモジュールのグラフ構造を読み解くことはできないようになっています。(GitHub Issueより)

ざっくりとしたイメージではモジュールはクラスのようなものと考えることができ、インプットを与えれば求めた処理を行い、アウトプットを返却するものです。モジュール内部のグラフなどを考える必要はありません。

f:id:kamujun:20180810161928p:plain

TensorFlow Hubのモジュールの用語・機能について説明します。

  • シグネチャ
    モジュールが持つメソッドのような機能です。モジュールは複数のシグネチャを持つことができ、指定したシグネチャによってインプットや処理内容などが変わります。指定しない場合はdefaultが指定されます。

  • アウトプット
    複数のアウトプットを持つモジュールがあります。何も指定しない場合はdefaultのアウトプットのみが返却されますが、as_dict＝Trueを指定することで定められたoutputがリストで返却されます。

提供モジュール

TensorFlow Hubでは様々な分野の事前学習済みモデルが用意されています。同一のモデルであっても返却される次元数や学習させる言語について異なるモジュールが用意されている場合があります。2018年8月現在では以下の通りですが、今後映像や音声系のモジュールの提供も予定しているとのこと。

画像系

  • Inception (V1, V2, V3)
  • Inception-ResNet (V2)
  • MobileNet (V1, V1 with TF-Lite, V2。それぞれmultiplierが25,50,75,100%のものを用意。)
  • NASNet-A (large, mobile)
  • PNASNet-5 (large)
  • ResNet (V1, V2)

テキスト系

  • NNLM (Chinese, English, German,
 Indonesian, Japanese, Korean, Spanish)
  • Universal Sentence Encoder (default, large, lite)
  • word2vec
  • ELMo

その他

  • Generative Adversarial Networks
    progan-128
  • DEep Local Features (DELF)
    delf
  • Inflated 3D Convnet (I3D)
    i3d-kinetics-400 i3d-kinetics-600

実装

次にTensorFlow Hubを用いた実装例を紹介します。利用するのは難しくないのですが、モジュールの作成は少し面倒なところがあります。以下で説明する機能の実装例はGitHubにあります。

モジュール利用・再学習

以下は公式サイトに記載されているサンプルコードです。前回の記事で紹介したELMoを例に紹介します。与えた文について単語ごとにELMoの分散表現を得る事を目的とした例です。

import tensorflow as tf
import tensorflow_hub as hub

with tf.Graph().as_default():
    elmo = hub.Module("https://tfhub.dev/google/elmo/2", trainable=True) # モジュール用意
    embeddings = elmo(
        ["the cat is on the mat", "dogs are in the fog"],
        signature="default",
        as_dict=True)["elmo"] # モジュール適用

    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        sess.run(tf.tables_initializer())

        print(sess.run(embeddings)) # 実行

モジュール用意

TensorFlow Hubを利用するにはモジュールのURLを指定してモジュールのダウンロードし、インスタンス生成を行います。その際にFine-tuning行わせたい場合は"trainable"をTrueにしておきます。これだけでグラフと重みのロードが終わっています。

モジュール適用

インスタンスに対してモジュールで定められた形式で引数を与えます。Googleが提供する各モジュールは仕様が公開されており、今回のELMoの場合はこちらのページに記載されています。ここでは分散表現を得る対象の文、シグネチャ、返り値の形式を指定します。

ここではシグネチャを"default"にしています。仕様ではこのモジュールは2つのシグネチャを持っており、一文を与えて処理させる"default"と、トークナイズした単語配列を与えて処理させる"tokens"が用意されています。

アウトプットについては"elmo"を指定しています。これは"default"のアウトプットでは文中単語の分散表現の平均を返すようになっており、目的とした単語ごとの分散表現は"elmo"で返却されるためです。"default"以外のアウトプットを用いるので"as_dict=True"としています。

実行

セッションで実行すればOKです。もちろん他の計算を足すことも可能です。Fine-tuningしたい場合は通常通りloss関数やオプティマイザを設定する必要があります。

モジュール作成

作成についてはコードを書くと長くなるので大まかな流れを図とともに説明します。本記事では概要の説明だけを行いますが、実装はGitHub上にあります。

モジュール作成で特にポイントとなるのはグラフの定義と重みを分けて扱う必要がある点です。これはエクスポート時にグラフ定義の関数と重みの割当は別々に行う必要があるためです。以下説明で順に説明していきます。

f:id:kamujun:20180810171347p:plain

①作成モデル準備

作成したいモデルを構築し、学習を行わせます。目標の性能が発揮することが確認できたら一度モデルを保存します。 もし重みをFine-tuning(再学習)させる想定が無いのであれば、convert_variables_to_constantsを用いてVariableをconstantに変換させておくと後の手順が一部スキップできます。

②モジュールスペック定義

モジュール内のグラフである"Module_Spec"を定義する関数を作成します。「①作成モデル準備」で保存したモデルをロードし、インプットとアウトプットのテンソルの変数を取得します。取得したインプットとアウトプットを用いてモジュールのシグネチャを追加することでスペック定義に用いる関数を用意します。

③モジュール作成

スペック定義を用いてモジュールのインスタンス生成を行います。この時点で作成したモデルと同じように利用することができますが、この時点ではVariableの値つまり重みは初期化されたままです。

④重み割り当て

「①作成モデル準備」で保存しておいたモデルの重みをロードして割り当てることで、事前学習した状態のモジュールとすることができます。ただし、Fine-tuningさせる予定がなく、constantに変換済みであればここの手順は不要です。(重みはグラフ中のconstantとして既に含まれているため)

⑤エクスポート

最後にこのモジュールを"export"したら事前学習済みモデルのモジュール完成です。指定したフォルダにモジュール一式であるpbファイルと重みであるVariableの値が出力されます。

自作モジュールの利用

作成したモデルは2つの方法で別環境で利用することができます。

モジュールファイルを直接受け渡し

exportしたモジュールを利用したい環境へ配置し、モジュールのダウンロードURLの部分でディレクトリを指定すれば利用かのうです。作成したモジュールが以下通りexportされた場合、ディレクトリmnist_moduleを指定すればOKです。

 .
 └mnist_module
    ├── assets
    ├── saved_model.pb
    ├── tfhub_module.pb
    └── variables

クラウドストレージを用いたホスティング

モジュールを圧縮した上でクラウドストレージに配置したものを利用することが可能です。その際は「.tar.gz」に圧縮すれば良いと説明があります。圧縮方法は公式の説明に任せますが、利用する際はダウンロードURLにクエリ「?tf-hub-format=compressed」を追加する必要があります。

mnist_module = hub.Module('https://your_hosting_address/mnist_module.tgz?tf-hub-format=compressed')

Google Cloud Storage と AWS S3 で試した結果、両者とも問題無く使用することができました。

TensorFlow Hubを使用した感想

実際に使用している中で感じた良い点とあと一歩の点を挙げて見ようと思います。

良い点

・大量リソースで学習された事前学習済みモデルが無料で利用出来る。しかも簡単に。
・各分野で有名なモデルを多数用意
・モデルによっては前処理まで含まれている（ボキャブラリ作成やトークナイズも含まれるので生の文を投げるだけで良い！）

難しい点

・KerasのLambdaレイヤにおいてTensorFlow hubを用いた場合、Fine−Tuningさせた結果の保存が不可能？
生のTenforFlowでは再学習結果も含めモデルを保存できたが、KerasのLambdaレイヤ内の重みの保存は実験する限りでは無理そうだった。カスタムレイヤで実装するしか無いかもしれない。

・モジュールの作成がやや煩雑
保存してあるモデルをサクッと読ませるだけでモジュール作成出来るようになると便利だと思います。

TensorFlow Hubによって事前学習済みモデルの恩恵を手軽に受けれるため、自分が取り組んでいるタスクに対して注力できるのはとても良いと感じました。有益な事前学習済みモデルを自作した場合にも他者が利用しやすい形で公開できるのは魅力ですね。