TensorFlow : Guide : Accelerators : TPU を利用する (翻訳/解説)

翻訳 : (株)クラスキャットセールスインフォメーション
更新日時 : 07/14/2018
作成日時 : 04/23/2018

* TensorFlow 1.9 でドキュメント構成が変わりましたので調整しました。
* 本ページは、TensorFlow 本家サイトの Guide – Accelerators – Using TPUs を翻訳した上で
適宜、補足説明したものです：

https://www.tensorflow.org/guide/using_tpu

* サンプルコードの動作確認はしておりますが、適宜、追加改変している場合もあります。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。

序文

このドキュメントは Cloud TPU を効果的に活用するために必要な主要な TensorFlow API をウォークスルーし、通常の TensorFlow 使用方法と TPU の使用方法の間の違いをハイライトします。

このドキュメントは次のようなユーザを対象にしています :

TensorFlow の Estimator と Dataset API に精通している。
既存のモデルを使用して多分 Cloud TPU を試したことがある。
サンプル TPU モデル [1] [2] (訳注: リンク切れ) のコードを多分大まかに読んだことがある。
既存の Estimator モデルを Cloud TPU 上で実行するためにポートすることに興味がある。

TPUEstimator

Estimator は TensorFlow のモデル・レベルの抽象です。標準的な Estimator は CPU と GPU 上でモデルを駆動できます。TPU 上でモデルを駆動するためには tf.contrib.tpu.TPUEstimator を使用しなければなりません。

TPUEstimator クラスは Estimator クラスとは幾分異なります。

CPU/GPU の両者あるいは Cloud TPU 上で動作可能なモデルを維持する最も単純な方法はモデルの推論フェイズ (入力から予測) を model_fn の外側で定義することです。それから Estimator setup と moderl_fn の別の実装を維持します、両者はこの推論ステップをラップします。このパターンのサンプルのために tensorflow/models の mnist.py と mnist_tpu.py 実装を比較します。

TPUEstimator をローカルで実行する

標準的な Estimator を作成するためには、コンストラクタを呼び出し、それに model_fn に渡します、例えば :

my_estimator = tf.estimator.Estimator(
  model_fn=my_model_fn)

貴方のローカルマシン上で tf.contrib.tpu.TPUEstimator を使用するために必要な変更は比較的少ないです。コンストラクタは２つの追加引数を要求します。use_tpu 引数を False に設定して、 config 引数として tf.contrib.tpu.RunConfig を渡すべきです、下で示されるように :

my_tpu_estimator = tf.contrib.tpu.TPUEstimator(
    model_fn=my_model_fn,
    config=tf.contrib.tpu.RunConfig()
    use_tpu=False)

丁度この単純な変更が TPUEstimator をローカルで実行することを可能にします。コマンドライン・フラグを次のように設定することにより、サンプル TPU モデルの大半はこのローカルモードで実行可能です :

$> python mnist_tpu.py --use_tpu=false --master=''

Note: この use_tpu=False 引数は TPUEstimator API を試すために有用です。それは完全な TPU 互換テストであることを意味しません。TPUEstimator のモデルのローカルでの成功的に実行することはそれが TPU 上で動作することを保証しません。

tpu.RunConfig を構築する

ローカル訓練のためにはデフォルト RunConfig は十分である一方で、これらの設定は実際の利用では無視できません。

Cloud TPU を使用するために切り替え可能な RunConfig のためのより典型的なセットアップは次のようなものです :

import tempfile
import subprocess

class FLAGS(object):
  use_tpu=False
  tpu_name=None
  # Use a local temporary path for the `model_dir`
  model_dir = tempfile.mkdtemp()
  # Number of training steps to run on the Cloud TPU before returning control.
  iterations = 50
  # A single Cloud TPU has 8 shards.
  num_shards = 8

if FLAGS.use_tpu:
    my_project_name = subprocess.check_output([
        'gcloud','config','get-value','project'])
    my_zone = subprocess.check_output([
        'gcloud','config','get-value','compute/zone'])
    cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(
            tpu_names=[FLAGS.tpu_name],
            zone=my_zone,
            project=my_project)
    master = tpu_cluster_resolver.get_master()
else:
    master = ''

my_tpu_run_config = tf.contrib.tpu.RunConfig(
    master=master,
    evaluation_master=master,
    model_dir=FLAGS.model_dir,
    session_config=tf.ConfigProto(
        allow_soft_placement=True, log_device_placement=True),
    tpu_config=tf.contrib.tpu.TPUConfig(FLAGS.iterations,
                                        FLAGS.num_shards),
)

それから tf.contrib.tpu.RunConfig をコンストラクタに渡さなければなりません :

my_tpu_estimator = tf.contrib.tpu.TPUEstimator(
    model_fn=my_model_fn,
    config = my_tpu_run_config,
    use_tpu=FLAGS.use_tpu)

典型的には FLAGS はコマンドライン引数で設定されます。ローカルでの訓練から Cloud TPU 上の訓練に切り替えるためには次が必要です :

FLAGS.use_tpu を True に設定する。
FLAGS.tpu_name を設定します、tf.contrib.cluster_resolver.TPUClusterResolver がそれを見つけられるようにです。
FLAGS.model_dir を Google Cloud Storage bucket url (gs://) に設定します。

Optimizer

cloud TPU 上で訓練するとき、optimizer を tf.contrib.tpu.CrossShardOptimizer でラップしなければなりません、これは
勾配を集めて結果を各シャード (各 TPU コア) にブロードキャストするために allreduce (訳注: 集団通信アルゴリズムの一つ) を使用します。

CrossShardOptimizer はローカル訓練と互換ではありません。そのため、同じコードをローカルと Cloud TPU 上の両者で実行するためには、次のような行を追加します :

optimizer = tf.train.GradientDescentOptimizer(learning_rate=learning_rate)
if FLAGS.use_tpu:
  optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer)

貴方のモデル・コードでグローバル FLAGS 変数を回避することを好む場合に、一つのアプローチは optimizer を Estimator の params の一つとして設定することです :

my_tpu_estimator = tf.contrib.tpu.TPUEstimator(
    model_fn=my_model_fn,
    config = my_tpu_run_config,
    use_tpu=FLAGS.use_tpu,
    params={'optimizer':optimizer})

Model 関数

このセクションは貴方が model 関数 (model_fn()) に (それを) TPUEstimator と互換にするために行わなければならない変更について詳述します。　

静的 shape

通常の使用の間は TensorFlow は各 tf.Tensor の shape をグラフ構築の間に決定することを試みます。実行の間任意の未知の shape 次元は動的に決定されます。

Cloud TPU 上で実行するために TensorFlow モデルは XLA を使用してコンパイルされます。XLA はコンパイル時に shape を決定するために同様のシステムを使用します。XLA は総ての tensor 次元がコンパイル時に静的に定義されることを要求します。総ての shape は定数として評価され、そして外部データ、あるいは変数やランダム数ジェネレータのようなステートフルな演算に依拠してはなりません。

Summaries

貴方のモデルから tf.summary の任意の使用を除去します。

TensorBoard summaries は貴方のモデルの内側を見る偉大な方法です。基本的な summaries の最小セットは TPUEstimator により model_dir のイベントファイルに自動的に記録されます。けれども、Cloud TPU 上で訓練するときカスタム summaries は現在サポートされていません。そのため TPUEstimator が依然としてローカルで summaries とともに動作しても、TPU 上で使用されればそれは失敗するでしょう。

Metrics

stand-alone metric_fn で貴方の評価メトリクス辞書を構築します。

評価メトリクスはモデル訓練の本質的なパートです。これらは Cloud TPU 上で完全にサポートされますが、僅かばかり異なるシンタクスを伴います。

標準的な tf.metrics は２つの tensor を返します。最初はメトリック値の実行平均を返し、その一方で２番目は実行平均を更新してこのバッチのための値を返します :

running_average, current_batch = tf.metrics.accuracy(labels, predictions)

標準的な Estimator ではこれらのペアの辞書を作成して、それを EstimatorSpec の一部として返します。

my_metrics = {'accuracy': tf.metrics.accuracy(labels, predictions)}

return tf.estimator.EstimatorSpec(
  ...
  eval_metric_ops=my_metrics
)

TPUEstimator では、代わりに関数 (これはメトリクス辞書を返します) と引数 tensor のリストを渡します、以下で示されるようにです :

def my_metric_fn(labels, predictions):
   return {'accuracy': tf.metrics.accuracy(labels, predictions)}

return tf.contrib.tpu.TPUEstimatorSpec(
  ...
  eval_metrics=(my_metric_fn, [labels, predictions])
)

TPUEstimatorSpec を使用する

TPUEstimatorSpec は hooks をサポートしません、そして幾つかのフィールドについては関数ラッパーが必要です。

Estimator の model_fn は EstimatorSpec を返さなければなりません。EstimatorSpec は、Estimator が相互作用する必要があるかもしれないモデルの総ての tf.Tensor を含む名前付けされたフィールドの単純な構造です。

TPUEstimators は tf.contrib.tpu.TPUEstimatorSpec を使用します。それと標準的な tf.estimator.EstimatorSpec の間には 2, 3 の違いがあります :

eval_metric_ops は metrics_fn にラップされなければなりません、このフィールドは eval_metrics に名前変更されます (上を参照)。
hooks は非サポートですので、これらのフィールドは省略されます。
もし使用されるのであれば、scaffold もまた関数にラップされなければなりません。このフィールドは scaffold_fn に名前変更されます。

Scaffold と Hooks はより進んだ利用のためで、典型的には省略可能です。

入力関数

入力関数はそれらが Cloud TPU 自身ではなく、ホストコンピュータで動作するとき殆ど変更なしで動作します。このセクションは２つの必要な調整を説明します。

Params 引数

標準的な Estimator のための input_fn は params 引数を含むことができます ; TPUEstimator のための input_fn は params 引数を含まなければなりません。これは、estimator に入力ストリームの各レプリカのためのバッチサイズを設定することを可能にするために必要です。TPUEstimator のための input_fn のための最小限のシグネチャは :

def my_input_fn(params):
  pass

ここで params[‘batch-size’] はバッチサイズを含みます。

静的 shape とバッチサイズ

貴方の input_fn により生成された入力パイプラインは CPU で動作します。そのためそれは大抵は XLA/TPU 環境により課せられる自由で厳密な静的 shape 要件です。一つの要件は貴方の入力パイプラインから TPU に供給されるデータのバッチは、標準的な TensorFlow shape 推論アルゴリズムにより決定される静的 shape を持つことです。中間的な tensor は動的 shape を自由に持ちます。もし shape 推論が失敗して、しかし shape が知られている場合には、tf.set_shape() を使用して正しい shape を課すことができます。

下のサンプルでは shape 推論アルゴリズムが失敗していますが、set_shape を使用して正されています :

>>> x = tf.zeros(tf.constant([1,2,3])+1)
>>> x.shape

TensorShape([Dimension(None), Dimension(None), Dimension(None)])

>>> x.set_shape([2,3,4])

多くの場合バッチサイズは唯一の未知の次元です。

tf.data を使用する、典型的な入力パイプラインは通常な固定サイズのバッチを生成します。けれども、有限の Dataset の最後のバッチは典型的にはより小さく、丁度残りの要素を含みます。Dataset はそれ自身の長さや有限性は知りませんので、標準的な batch メソッドは総てのバッチがそれ自身の上で固定サイズ・バッチを持つかは決定できません :

>>> params = {'batch_size':32}
>>> ds = tf.data.Dataset.from_tensors([0, 1, 2])
>>> ds = ds.repeat().batch(params['batch-size'])
>>> ds

<BatchDataset shapes: (?, 3), types: tf.int32>

最も簡単な fix は次のように tf.contrib.data.batch_and_drop_remainder を apply することです :

>>> params = {'batch_size':32}
>>> ds = tf.data.Dataset.from_tensors([0, 1, 2])
>>> ds = ds.repeat().apply(
...     tf.contrib.data.batch_and_drop_remainder(params['batch-size']))
>>> ds

 <_RestructuredDataset shapes: (32, 3), types: tf.int32>

このアプローチの一つの欠点は、その名前が意味するように、このバッチ処理 method は dataset の最後で任意の端数のバッチを投げ捨てます。これは訓練のために使用される dataset を無限に反復するためには良いのですが、もし貴方が正確な数のエポックのために訓練することを望む場合には問題になるかもしれません。正確に 1-エポックの評価を遂行するためには貴方自身の tf.metrics を作成するときにバッチの長さを手動でパディングし、パディング要素をゼロ重みを持つように設定することによりこれを回避することができます。

Dataset

Cloud TPU を使用するとき tf.data.Dataset API の効率的な使用は重要です、何故ならばデータを十分に迅速に供給できない場合には Cloud TPU を使用することは不可能だからです。dataset パフォーマンスの詳細については Input Pipeline Performance Guide を見てください。

(tf.data.Dataset.from_tensor_slice か他の in-graph データを使用する) 最も単純な実験を別にすれば TPUEstimator の Dataset により読まれるデータファイルの総ては Google Cloud Storage Bucket にストアする必要があるでしょう。

殆どのユースケースについて、貴方のデータを TFRecord フォーマットに変換してそれを読むために tf.data.TFRecordDataset を使用することを推奨します。けれどもこれは厳しい要件ではなくて好ましいのであれば他の dataset リーダー (FixedLengthRecordDataset か TextLineDataset) を使用できます。

tf.data.Dataset.cache を使用して小さい dataset は全体をメモリにロードすることができます。

使用されるデータ・フォーマットにかかわらず、100 MB のオーダー上の巨大なファイルを使用することが強く推奨されます。これはこのネットワーク化された設定において特に重要です、というのはファイルを開くオーバーヘッドが本質的により高いからです。

使用されるリーダーのタイプにかかわらず、コンストラクタへの buffer_size 引数を使用してバッファリングを有効にすることもまた重要です。この引数はバイトで指定されます。必要なときにデータが利用可能であるように最小限 2, 3 MB (buffer_size=8*1024*1024) が推奨されます。

TPU-demos レポジトリは imagenet データセットをダウンロードしてそれを適切なフォーマットに変換するためのスクリプト (訳注: リンク切れ) を含みます。これはレポジトリに含まれる imagenet モデル (訳注: リンク切れ) と一緒にこれらのベストプラクティスを示します。

以上

月	火	水	木	金	土	日
						1
2	3	4	5	6	7	8
9	10	11	12	13	14	15
16	17	18	19	20	21	22
23	24	25	26	27	28	29
30