TensorFlow Hub を Cloud Dataflowで使ってみた

Published in

google-cloud-jp

8 min readApr 6, 2018

先日のTF Dev Summit 2018で発表されたTensorFlowの新機能に、TensorFlow Hubというのがあります。これはDockerでいえばDockerHubに該当するもので、モデルを簡単に共有できるエコシステムの一部です。本記事執筆時点(2018年4月)ではまだ一般ユーザーが独自のモデルをHubへアップロードすることはできませんが、現在多くの有用な学習済みモデルがGoogleから提供されています。

何の記事？
上記のTensorFlow HubをDataflowから使い、バッチでたくさんの画像の特徴量を抽出してみよう、というお話です。

Dataflowって何？
Cloud DataflowとはGoogle Cloud Platformのプロダクトの1つで、いわゆるETLに該当します。サーバーレス環境なので、ほんの数十行のコードでオートスケールする大規模なデータ処理が可能です。実はこのDataflow、TensorFlowとも非常に相性がよく、TFRecord形式で出力もできるので学習データの前処理なんかもできてしまいます！

TensorFlow Hubの使い方

TensorFlow Hubを使うのはとても簡単です。まずはライブラリのインストールをします。TensorFlowのバージョンは1.7以降が必須なので、すでに下位バージョンをインストール済みの環境では-U をつけてアップグレードしましょう。

pip install “tensorflow>=1.7.0”
pip install tensorflow-hub

次にTF Hubで学習済みモデルをロードします。

import tensorflow as tf
import tensorflow_hub as hub
module = hub.Module(“https://tfhub.dev/google/imagenet/inception_v3/feature_vector/1")

以上！簡単ですね。これはちなみにInception-v3というモデルにImagenetの画像データを学習させたものです。通常は1000クラスを識別するモデルですが、TopのDenseレイヤーがすでに切り離されている状態で用意されており、新たなDenseレイヤーを追加して転移学習がしやすくなっています。

DataflowでTensorFlow Hubを使うには

DataflowにもTensorFlowはプリインストール済みですが、現時点ではv1.7が入っていないので、別途インストールしてあげる必要があります。

Dataflowにライブラリを追加する方法はいくつかありますが、今回はsetup.pyを使う方法を紹介します。まずは次のようなsetup.pyを作成します。

import setuptoolssetuptools.setup(
 name=’df-tfhub’,
 version=’0.1',
 install_requires=[‘tensorflow>=1.7.0’, ‘tensorflow-hub’],
 packages=setuptools.find_packages(),
)

次にsetup.pyをPipelineOptionsに設定します。

options = beam.options.pipeline_options.PipelineOptions()
setup_options = options.view_as(beam.options.pipeline_options.SetupOptions)
setup_options.setup_file = './setup.py'

これで終わりです。あとは自由にTF Hubをつかえます。

パイプラインの定義

Dataflowは行いたい処理を数珠つなぎに繋げていって、それらをパイプラインで処理することができます。今回は上記のInception-v3モデルを使って、GCS上にある画像の特徴量を出力するパイプラインを構築してみましょう。次の図のようなパイプラインです。

それぞれ順番に解説すると、

画像ファイルがあるGCSのパスを入力（固定）
パス内のファイル一覧を取得
画像をロード（＆リサイズ）
特徴量抽出（TF Hubを使っているところ）
GCSに書き込み

という感じです。Dataflowでこれらのパイプラインを定義するには、次のようになります。

p = beam.Pipeline(options=options)
(p | ‘gcs_path’ >> beam.Create([‘gs://[BUCKET]’])
   | ‘get file names’ >> beam.FlatMap(get_file_names)
   | ‘load image’ >> beam.Map(load_image)
   | ‘extract features’ >> beam.Map(extract_features)
   | ‘write’ >> beam.io.WriteToText(‘gs://[/BUCKET]/hub.txt’,
                                   num_shards=1)
)

Dataflow(Python)では、パイプラインの処理を繋げるにはパイプ「 | 」で繋げるということをします。通常のPythonコードとは感覚が違いますが、こういうものだと思いましょう。

beam.Map() は、1つの入力に対して1つの出力をします。beam.FlatMap()は、1つの入力に対して複数の出力をします。ここでは、GCSからファイル名を取得している処理をFlatMapにし、それ以外はMapで１画像を１処理単位としています。

beam.Map()とbeam.FlatMap()の引数は、それぞれ具体的にその中で行なっている処理の関数です。次のように定義します。

# ファイル名取得
def get_file_names(gcs_path):
  import tensorflow as tf
 
  for fn in tf.gfile.ListDirectory(gcs_path):
    yield gcs_path + ‘/’ + fn# 画像をGCSからロード
def load_image(fn):
  import tensorflow as tf
  import numpy as np
  from PIL import Image
 
  f = tf.gfile.Open(fn)
  im = Image.open(f)
  im = im.resize((299, 299))
 
  return np.array(im).reshape((1, 299, 299, 3))# Hubで読み込んだInception-v3で特徴量抽出
def extract_features(elem):
  import tensorflow as tf
  import tensorflow_hub as hub
 
  with tf.Graph().as_default():
    module = hub.Module(“https://tfhub.dev/google/imagenet/inception_v3/feature_vector/1")
    features = module(elem)
  with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    return sess.run(features)

それぞれ、数行のコードで済んでしまうのがDataflowのすごいところですね！ちなみにDataflowの処理関数（Transform）上からGCSへアクセスする方法はいくつかありますが、tf.gfileを使うと簡単です。

Takeaways..

Dataflowは冒頭で記載したように、ワークロードに合わせて自動でスケールしてくれます。ただし、上記のサンプルコードはちょっと非効率な部分が含まれています。

ライブラリのロード
各関数内でライブラリをインポートしているのに気が付いたかもしれません。Dataflowでは、複数のワーカーで並列処理を行うため、どの処理をどのワーカーが実行しているのかわかりません。なので、関数内で使うライブラリは、関数内でインポートしないといけないのです。でもこれって結構オーバーヘッドにならないでしょうか？

解決方法の1つは、よく使うライブラリ（とTF Hubでロードしたモデル）はSingletonで定義してしまうことです。そうすれば、少なくとも１要素の処理単位でロードされてしまうことは防げます。他にも、いろんな方法があるので、みなさんも是非Dataflowを利用し、情報をシェアしてくださいね！

Disclaimer この記事は個人的なものです｡ここで述べられていることは私の個人的な意見に基づくものであり､私の雇用者には関係はありません｡

TensorFlow Hub を Cloud Dataflowで使ってみた

TensorFlow Hubの使い方

DataflowでTensorFlow Hubを使うには

パイプラインの定義

Takeaways..

Written by Hayato Yoshikawa