このドキュメントでは、Inception モデルの側面と、このモデルが Cloud TPU で効率的に動作するようにこれらの側面がどのように連携するかについて説明します。これは、Cloud TPU で Inception v3 を実行するための詳細なガイドです。大幅な改善につながったモデルの特定の変更について詳細に説明します。このドキュメントは、Inception v3 のチュートリアルを補足するものです。
Inception v3 の TPU トレーニングの実行は、同様の構成の GPU ジョブによって生成された精度曲線と一致します。このモデルは、v2-8、v2-128、v2-512 の構成でトレーニングを完了しています。このモデルは、それぞれについて約 170 エポックで 78.1% を超える精度を達成しています。
このドキュメントに示されているコードサンプルは、実際の実装で何が起こるかを簡単に説明するためのものです。実際のコードは GitHub をご覧ください。
はじめに
Inception v3 は、広く使用されている画像認識モデルで、ImageNet データセットで 78.1% を超える精度を達成しています。このモデルは、長年にわたって複数の研究者によって開発された多くのアイデアが結実したものです。これは、Szegedy 氏他の『Rethinking the Inception Architecture for Computer Vision』という論文をベースにしています。
このモデル自体は、畳み込み層、平均プーリング層、最大プーリング層、連結層、ドロップアウト層、全結合層など、対称と非対称の構成要素で構成されています。Batchnorm はモデル全体で広く使用され、活性化入力に適用されます。損失は Softmax によって計算されます。
モデルの概要図を以下に示します。
Inception アーキテクチャの詳細については、Inception モデルの README をご覧ください。
Estimator API
Inception v3 の TPU バージョンは、開発を容易にするために設計された API である TPUEstimator を使用して作成されているため、基盤となるハードウェアの詳細ではなく、モデルそのものに集中できます。この API は、チェックポイントの保存や復元などの一般的な機能を自動化しながら、バックグラウンドで TPU でモデルを実行するために必要な低水準の重作業のほとんどを行います。
Estimator API は、コードのモデル部分と入力部分を分けます。モデル定義と TensorFlow グラフの入力パイプライン / 前処理ステージにそれぞれ対応する、model_fn 関数と input_fn 関数を定義する必要があります。以下は、これらの関数のサンプルです。
def model_fn(features, labels, mode, params):
…
return tpu_estimator.TPUEstimatorSpec(mode=mode, loss=loss, train_op=train_op)
def input_fn(params):
def parser(serialized_example):
…
return image, label
…
images, labels = dataset.make_one_shot_iterator().get_next()
return images, labels
API によって提供される 2 つの重要な関数は train() と evaluate() で、それぞれトレーニングと評価に使用されます。これらは通常、main 関数内で呼び出されます。以下にこの例を示します。
def main(unused_argv):
…
run_config = tpu_config.RunConfig(
master=FLAGS.master,
model_dir=FLAGS.model_dir,
session_config=tf.ConfigProto(
allow_soft_placement=True, log_device_placement=True),
tpu_config=tpu_config.TPUConfig(FLAGS.iterations, FLAGS.num_shards),)
estimator = tpu_estimator.TPUEstimator(
model_fn=model_fn,
use_tpu=FLAGS.use_tpu,
train_batch_size=FLAGS.batch_size,
eval_batch_size=FLAGS.batch_size,
config=run_config)
estimator.train(input_fn=input_fn, max_steps=FLAGS.train_steps)
eval_results = inception_classifier.evaluate(
input_fn=imagenet_eval.input_fn, steps=eval_steps)
ImageNet データセット
画像認識にモデルを使用するには、その前にモデルをトレーニングする必要があります。これは通常、ラベル付けされた大規模な画像セットを使用して教師あり学習によって行われます。Inception v3 はさまざまなラベル付き画像セットを使用してトレーニングできますが、ImageNet は最適な一般のデータセットです。
ImageNet には、1,000 万を超える、ラベル付き画像の URL があります。約 100 万枚の画像には、ラベル付きオブジェクトのより正確な位置を指定する境界ボックスもあります。
このモデルの場合、ImageNet データセットは 1,331,167 枚の画像で構成され、1,281,167 枚の画像を含むトレーニング用データセットと 50,000 枚の画像を含む評価用のデータセットに分割されています。
トレーニング用と評価用のデータセットは意図的に分割されています。モデルのトレーニングにはトレーニング用データセットの画像のみが使用され、モデルの精度の評価には評価用データセットの画像のみが使用されます。
このモデルは、画像を TFRecord として保存することを想定しています。画像を元の JPEG ファイルから TFRecord に変換するには、オープンソースのバッチ スクリプト download_and_preprocess_imagenet.sh を使用します。このスクリプトにより、以下の形式の一連のファイル（トレーニング用と検証用の両方）が生成されます。
${DATA_DIR}/train-00000-of-01024
${DATA_DIR}/train-00001-of-01024
...
${DATA_DIR}/train-01023-of-01024
and
${DATA_DIR}/validation-00000-of-00128
S{DATA_DIR}/validation-00001-of-00128
...
${DATA_DIR}/validation-00127-of-00128
DATA_DIR は、セットがある場所です（例: DATA_DIR=$HOME/imagenet-data）。
Inception モデルの README の「Getting Started」セクションには、このスクリプトの作成および実行方法の詳細な手順が記載されています。
入力パイプライン
各 Cloud TPU デバイスには 8 つのコアがあり、ホスト（CPU）に接続されています。より大きいスライスには複数のホストがあります。他のより大規模な構成は、複数のホストと通信します。たとえば、v2-256 は 16 のホストと通信します。
ホストは、ファイル システムまたはローカルメモリからデータを取得し、必要なデータ前処理を行い、前処理されたデータを TPU のコアに転送します。ホストが個別に行うデータ処理の 3 つのフェーズ 1）保存、2）前処理、3）転送について考えます。概要を以下の図に示します。
良好なパフォーマンスを得るには、システムのバランスを取る必要があります。ホスト CPU が画像の取得、デコード、関連する前処理を行うのにかかる時間は、理想的には、TPU が計算を行うのにかかる時間よりわずかに短いか、ほぼ同じである必要があります。ホスト CPU が 3 つのデータ処理フェーズを完了するのに TPU よりも時間がかかる場合、実行はホストバウンドになります（注: TPU は非常に高速であるため、非常に単純なモデルではこれを回避できない場合があります）。両方のケースを以下の図に示します。
Inception v3 の現在の実装は、入力バウンドです。画像はファイル システムから取得し、デコードしてから前処理する必要があります。中程度のものから複雑なものまで、さまざまなタイプの前処理ステージを使用できます。最も複雑な前処理ステージを使用すると、前処理ステージで実行される多数の高負荷なオペレーションによりシステムの処理能力が低下し、トレーニング パイプラインは前処理バウンドになります。ただし、78.1% を超える精度を達成するためにそのレベルの複雑さを使用する必要はありません。代わりに、天秤のもう一方を重くして、モデルを TPU バウンドにする、中程度に複雑な前処理ステージを使用します。これについては、次のセクションで詳しく説明します。
このモデルは tf.data.Dataset を使用して、すべての入力パイプラインに関連するニーズを処理します。tf.data API を使用して入力パイプラインを最適化する方法の詳細については、データセットのパフォーマンス ガイドをご覧ください。
関数を定義して、それを Estimator API に渡すことはできますが、Inception の場合は、代わりに InputPipeline クラスを作成して必要なすべての機能をカプセル化し、__call__ メソッドを定義します。
Estimator API を使用すると、このクラスを使用するのが非常に簡単になります。以下のサンプルコード スニペットに示すように、関数 train() と evaluate() の input_fn パラメータにこのクラスを渡すだけです。
def main(unused_argv):
…
inception_classifier = tpu_estimator.TPUEstimator(
model_fn=inception_model_fn,
use_tpu=FLAGS.use_tpu,
config=run_config,
params=params,
train_batch_size=FLAGS.train_batch_size,
eval_batch_size=eval_batch_size,
batch_axis=(batch_axis, 0))
…
for cycle in range(FLAGS.train_steps // FLAGS.train_steps_per_eval):
tf.logging.info('Starting training cycle %d.' % cycle)
inception_classifier.train(
input_fn=InputPipeline(True), steps=FLAGS.train_steps_per_eval)
tf.logging.info('Starting evaluation cycle %d .' % cycle)
eval_results = inception_classifier.evaluate(
input_fn=InputPipeline(False), steps=eval_steps, hooks=eval_hooks)
tf.logging.info('Evaluation results: %s' % eval_results)
以下のコード スニペットに InputPipeline クラスの主な要素が表示されています。3 つのフェーズを異なる色で強調表示しています。メソッド __call__ は、tf.data.Dataset を使用してデータセットを作成し、データセットの組み込みプリフェッチ機能、インターリーブ機能、シャッフル機能を利用するために、一連の API 呼び出しを行います。
class InputPipeline(object):
def __init__(self, is_training):
self.is_training = is_training
def __call__(self, params):
# Storage
file_pattern = os.path.join(
FLAGS.data_dir, 'train-*' if self.is_training else 'validation-*')
dataset = tf.data.Dataset.list_files(file_pattern)
if self.is_training and FLAGS.initial_shuffle_buffer_size > 0:
dataset = dataset.shuffle(
buffer_size=FLAGS.initial_shuffle_buffer_size)
if self.is_training:
dataset = dataset.repeat()
def prefetch_dataset(filename):
dataset = tf.data.TFRecordDataset(
filename, buffer_size=FLAGS.prefetch_dataset_buffer_size)
return dataset
dataset = dataset.apply(
tf.contrib.data.parallel_interleave(
prefetch_dataset,
cycle_length=FLAGS.num_files_infeed,
sloppy=True))
if FLAGS.followup_shuffle_buffer_size > 0:
dataset = dataset.shuffle(
buffer_size=FLAGS.followup_shuffle_buffer_size)
# Preprocessing
dataset = dataset.map(
self.dataset_parser,
num_parallel_calls=FLAGS.num_parallel_calls)
dataset = dataset.prefetch(batch_size)
dataset = dataset.apply(
tf.contrib.data.batch_and_drop_remainder(batch_size))
dataset = dataset.prefetch(2) # Prefetch overlaps in-feed with training
images, labels = dataset.make_one_shot_iterator().get_next()
# Transfer
return images, labels
storage セクションは、データセットの作成から始まり、ストレージからの TFRecord の読み取りが含まれます（tf.data.TFRecordDataset を使用）。特殊用途の関数 repeat() と shuffle() は必要に応じて使用されます。関数 tf.contrib.data.parallel_interleave() は、その入力内で関数 prefetch_dataset() をマッピングして、ネストされたデータセットを生成し、インターリーブされたその要素を出力します。並行して、ネストされたデータセット cycle_length から要素を取得します。これにより、スループットが向上します。sloppy 引数は、出力が決定論的な順序で生成されるという要件を緩和し、リクエストされたときに要素をすぐに利用できないネストされたデータセットを実装がスキップできるようにします。
preprocessing セクションは、dataset.map(parser) を呼び出します。これは、画像が前処理されるパーサー関数を呼び出します。前処理ステージの詳細については、次のセクションで説明します。
transfer セクション（この関数の最後にある）には、return images, labels という行があります。TPUEstimator は返された値を受け取り、デバイスに自動的に転送します。
以下の図は、Inception v3 の Cloud TPU パフォーマンス トレースのサンプルを示しています。TPU の計算時間は、送り込み時間を差し引いて、現在 815 ミリ秒程度です。
ホスト ストレージもトレース上に表示されています。以下に示します。
画像のデコード機能と一連の画像歪み機能を含むホストの前処理を以下に示します。
ホスト / TPU の転送はここで確認できます。
前処理ステージ
画像の前処理はシステムの重要な部分であり、トレーニング中にモデルが達成する最高精度に大きく影響する可能性があります。少なくとも、画像をデコードし、モデルに合わせてサイズを変更する必要があります。Inception の場合、画像は 299x299x3 ピクセルである必要があります。
ただし、デコードとサイズ変更だけでは高い精度を得るには十分ではありません。ImageNet のトレーニング用データセットには、1,281,167 枚の画像が含まれています。トレーニング用画像セットを通過する 1 回のパスは、エポックと呼ばれます。トレーニング中、モデルは画像認識能力を向上させるためにトレーニング用データセットを複数回通過する必要があります。Inception v3 の場合、必要なエポック数は、グローバル バッチサイズに応じて 140〜200 の範囲内になります。
画像をモデルにフィードする前に画像を継続的に変更し、特定の画像がエポックごとにわずかに異なる方法で行うことは、非常に有益です。この画像の前処理を最適に行う方法は、科学的・論理的であると同時に職人技でもあります。一方で、よく設計された前処理ステージは、モデルの認識能力を大幅に向上させることができます。他方で、あまりにも単純な前処理ステージは、同じモデルがトレーニング中に達成できる最高精度が制限されてしまう可能性があります。
Inception v3 では、比較的単純で計算負荷が低いものから、かなり複雑で計算負荷が高いものまで、前処理ステージにさまざまなオプションが用意されています。そのような異なる 2 つのタイプは、ファイル vgg_preprocessing.py と inception_preprocessing.py にあります。
ファイル vgg_preprocessing.py は、resnet を 75% の精度までトレーニングするために使用された前処理ステージを定義しますが、Inception v3 に適用すると最適な結果を得られません。
ファイル inception_preprocessing.py には、TPU での実行時に Inception v3 を 78.1～78.5% の精度までトレーニングするために使用された、複雑さのレベルが異なる複数のオプションの前処理ステージが含まれています。このセクションでは、前処理パイプラインについて説明します。
前処理は、モデルがトレーニング中か、推論 / 評価に使用されているかによって異なります。
評価時の前処理は非常に簡単です。画像の中央部を切り抜き、デフォルトの 299x299 のサイズに変更します。これを行うコード スニペットを以下に示します。
def preprocess_for_eval(image, height, width, central_fraction=0.875):
with tf.name_scope(scope, 'eval_image', [image, height, width]):
if image.dtype != tf.float32:
image = tf.image.convert_image_dtype(image, dtype=tf.float32)
image = tf.image.central_crop(image, central_fraction=central_fraction)
image = tf.expand_dims(image, 0)
image = tf.image.resize_bilinear(image, [height, width], align_corners=False)
image = tf.squeeze(image, [0])
image = tf.subtract(image, 0.5)
image = tf.multiply(image, 2.0)
image.set_shape([height, width, 3])
return image
トレーニング中の切り抜きはランダムに行われます。境界ボックスがランダムに選択され、画像のリージョンが選択され、そのリージョンがサイズ変更されます。サイズ変更された画像は、必要に応じて反転され、その色が歪められます。これを行うコード スニペットを以下に示します。
def preprocess_for_train(image, height, width, bbox, fast_mode=True, scope=None):
with tf.name_scope(scope, 'distort_image', [image, height, width, bbox]):
if bbox is None:
bbox = tf.constant([0.0, 0.0, 1.0, 1.0], dtype=tf.float32, shape=[1, 1, 4])
if image.dtype != tf.float32:
image = tf.image.convert_image_dtype(image, dtype=tf.float32)
distorted_image, distorted_bbox = distorted_bounding_box_crop(image, bbox)
distorted_image.set_shape([None, None, 3])
num_resize_cases = 1 if fast_mode else 4
distorted_image = apply_with_random_selector(
distorted_image,
lambda x, method: tf.image.resize_images(x, [height, width], method),
num_cases=num_resize_cases)
distorted_image = tf.image.random_flip_left_right(distorted_image)
if FLAGS.use_fast_color_distort:
distorted_image = distort_color_fast(distorted_image)
else:
num_distort_cases = 1 if fast_mode else 4
distorted_image = apply_with_random_selector(
distorted_image,
lambda x, ordering: distort_color(x, ordering, fast_mode),
num_cases=num_distort_cases)
distorted_image = tf.subtract(distorted_image, 0.5)
distorted_image = tf.multiply(distorted_image, 2.0)
return distorted_image
関数 distort_color は色を変更します。これは明度と彩度のみが変更される高速モードを提供します。フルモードでは、明度、彩度、色相が変更され、これらが変更される順序がランダムに変更されます。この関数のコード スニペットを以下に示します。
def distort_color(image, color_ordering=0, fast_mode=True, scope=None):
with tf.name_scope(scope, 'distort_color', [image]):
if fast_mode:
if color_ordering == 0:
image = tf.image.random_brightness(image, max_delta=32. / 255.)
image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
else:
image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
image = tf.image.random_brightness(image, max_delta=32. / 255.)
else:
if color_ordering == 0:
image = tf.image.random_brightness(image, max_delta=32. / 255.)
image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
image = tf.image.random_hue(image, max_delta=0.2)
image = tf.image.random_contrast(image, lower=0.5, upper=1.5)
elif color_ordering == 1:
image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
image = tf.image.random_brightness(image, max_delta=32. / 255.)
image = tf.image.random_contrast(image, lower=0.5, upper=1.5)
image = tf.image.random_hue(image, max_delta=0.2)
elif color_ordering == 2:
image = tf.image.random_contrast(image, lower=0.5, upper=1.5)
image = tf.image.random_hue(image, max_delta=0.2)
image = tf.image.random_brightness(image, max_delta=32. / 255.)
image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
elif color_ordering == 3:
image = tf.image.random_hue(image, max_delta=0.2)
image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
image = tf.image.random_contrast(image, lower=0.5, upper=1.5)
image = tf.image.random_brightness(image, max_delta=32. / 255.)
return tf.clip_by_value(image, 0.0, 1.0)
関数 distort_color は計算コストが高くなります。これは、色相と彩度にアクセスするために必要な非線形の RGB から HSV への変換と HSV から RGB への変換に原因の一端があります。高速モードとフルモードの両方でこれらの変換が必要です。高速モードの方が計算コストは低いのですが、これを有効にした場合はモデルを CPU 計算バウンドリージョンに push します。
代わりに、新しい関数 distort_color_fast がオプションのリストに追加されました。この関数は、JPEG 変換スキームを使用して画像を RGB から YCrCb にマッピングし、明度と Cr / Cb 色度をランダムに変更してから RGB にマッピングし直します。この関数のコード スニペットを以下に示します。
def distort_color_fast(image, scope=None):
with tf.name_scope(scope, 'distort_color', [image]):
br_delta = random_ops.random_uniform([], -32./255., 32./255., seed=None)
cb_factor = random_ops.random_uniform(
[], -FLAGS.cb_distortion_range, FLAGS.cb_distortion_range, seed=None)
cr_factor = random_ops.random_uniform(
[], -FLAGS.cr_distortion_range, FLAGS.cr_distortion_range, seed=None)
channels = tf.split(axis=2, num_or_size_splits=3, value=image)
red_offset = 1.402 * cr_factor + br_delta
green_offset = -0.344136 * cb_factor - 0.714136 * cr_factor + br_delta
blue_offset = 1.772 * cb_factor + br_delta
channels[0] += red_offset
channels[1] += green_offset
channels[2] += blue_offset
image = tf.concat(axis=2, values=channels)
image = tf.clip_by_value(image, 0., 1.)
return image
以下は、前処理が行われたサンプル画像です。画像のランダムに選択されたリージョンが選択され、distort_color_fast 関数を使用して色が変更されました。
関数 distort_color_fast は計算効率がよく、さらにトレーニングを TPU 実行時間バウンドにすることができます。また、よい結果が得られ、これを使用して Inception v3 モデルをトレーニングし、1,024〜16,384 のバッチサイズで 78.1% 以上の精度を達成しました。Inception v3 のデフォルトとしてこれが使用されます。
オプティマイザー
現在のモデルでは、SGD、Momentum、RMSProp という 3 種類のオプティマイザーを紹介しています。
Stochastic gradient descent (SGD) は最も単純な種類の更新であり、重みは負の勾配方向に移動します。その単純さにもかかわらず、一部のモデルでは良好な結果が得られます。更新のダイナミクスは以下のように表すことができます。
Momentum は一般的なオプティマイザーであり、SGD が達成できるよりも早く収束することがあります。このオプティマイザーは、SGD と同様に重みを更新しますが、前の更新の方向にもコンポーネントを追加します。更新のダイナミクスは以下のようになります。
これは、以下のように表すことができます。
最後の項は、前の更新の方向のコンポーネントです。以下の図にこれを示します。
慣性 \({\beta}\) には、一般的に使用される値 0.9 を使用します。
RMSprop は、講義で Geoff Hinton が最初に提唱した一般的なオプティマイザーです。更新のダイナミクスは以下のようになります。
Inception v3 でのテストにより、RMSProp は最高精度とその到達時間において最良の結果を得られ、慣性は 1 秒に近いことが示されています。したがって、RMSprop がデフォルトのオプティマイザーとして設定されています。使用するパラメータは、減衰 \({\alpha}\) = 0.9、慣性 \({\beta}\) = 0.9、\({\epsilon}\) = 1.0 です。
オプティマイザーの選択肢を含むコード スニペットは以下のようになります。
if FLAGS.optimizer == 'sgd':
tf.logging.info('Using SGD optimizer')
optimizer = tf.train.GradientDescentOptimizer(
learning_rate=learning_rate)
elif FLAGS.optimizer == 'momentum':
tf.logging.info('Using Momentum optimizer')
optimizer = tf.train.MomentumOptimizer(
learning_rate=learning_rate, momentum=0.9)
elif FLAGS.optimizer == 'RMS':
tf.logging.info('Using RMS optimizer')
optimizer = tf.train.RMSPropOptimizer(
learning_rate,
RMSPROP_DECAY,
momentum=RMSPROP_MOMENTUM,
epsilon=RMSPROP_EPSILON)
else:
tf.logging.fatal('Unknown optimizer:', FLAGS.optimizer)
TPU で実行し、Estimator API を使用する場合は、オプティマイザーを CrossShardOptimizer 関数でラップして、（必要な相互通信とともに）レプリカ間の同期を確実に行うようにする必要があります。Inception v3 でこれを行うコード スニペットを以下に示します。
if FLAGS.use_tpu:
optimizer = tpu_optimizer.CrossShardOptimizer(optimizer)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
train_op = optimizer.minimize(loss, global_step=global_step)
指数移動平均
トレーニング中の通常の動作過程では、オプティマイザーの更新ルールに従って、トレーニング可能なパラメータが逆伝播中に更新されます。これらは前のセクションで説明しましたが、便宜上ここでも繰り返します。
指数移動平均（指数平滑とも呼ばれます）は、更新された重みに適用されるオプションの後処理ステップであり、パフォーマンスの著しい向上につながることがあります。Inception v3 では、この追加のステップを行うことにより、大きなメリットが得られます。TensorFlow は関数 tf.train.ExponentialMovingAverage を提供し、以下の数式を使用して、重み \({\theta}\) の指数移動平均 \({\hat{\theta}}\) を計算します。
\({\alpha}\) は減衰係数（1.0 に近い）です。Inception v3 の場合、\({\alpha}\) は 0.995 に設定されます。
これは無限インパルスレスポンス（IIR）フィルタですが、減衰係数は、次の図に示すように、ほとんどのエネルギー（関連するサンプル）が存在する有効なウィンドウを確立します。
これをより明確にするために、次のようにフィルタ式を書き換えます。
ここでは \({\hat\theta_{-1}}=0\) を使用しました。
\({\alpha}^k\) の値は k が大きくなるにつれ減衰するので、事実上サンプルのサブセットのみが \(\hat{\theta}_{t+T+1}\) に大きな影響を与えます。上記ウィンドウの期間の経験則は、\(\frac {1} {1-\alpha}\) になります。= 0.995 の場合は \({\alpha}\) = 200 に相当します。
最初にトレーニング可能な変数のコレクションを取得し、apply() メソッドを使用してトレーニングされた各変数のシャドウ変数を作成します（対応する演算を追加してシャドウコピーでこれらの移動平均を維持します）。Inception v3 でこれを行うコード スニペットを以下に示します。
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
train_op = optimizer.minimize(loss, global_step=global_step)
if FLAGS.moving_average:
ema = tf.train.ExponentialMovingAverage(
decay=MOVING_AVERAGE_DECAY, num_updates=global_step)
variables_to_average = (tf.trainable_variables() +
tf.moving_average_variables())
with tf.control_dependencies([train_op]), tf.name_scope('moving_average'):
train_op = ema.apply(variables_to_average)
評価中に指数移動平均変数を使用します。これを行うには、メソッド variables_to_restore() をチェックポイント ファイルに適用して、シャドウ変数名を使用して評価するクラス LoadEMAHook を定義します。
class LoadEMAHook(tf.train.SessionRunHook):
def __init__(self, model_dir):
super(LoadEMAHook, self).__init__()
self._model_dir = model_dir
def begin(self):
ema = tf.train.ExponentialMovingAverage(MOVING_AVERAGE_DECAY)
variables_to_restore = ema.variables_to_restore()
self._load_ema = tf.contrib.framework.assign_from_checkpoint_fn(
tf.train.latest_checkpoint(self._model_dir), variables_to_restore)
def after_create_session(self, sess, coord):
tf.logging.info('Reloading EMA...')
self._load_ema(sess)
hooks 関数は、以下のスニペットに示すように evaluate() に渡されます。
if FLAGS.moving_average:
eval_hooks = [LoadEMAHook(FLAGS.model_dir)]
else:
eval_hooks = []
…
eval_results = inception_classifier.evaluate(
input_fn=InputPipeline(False), steps=eval_steps, hooks=eval_hooks)
バッチ正規化
バッチ正規化は、収束時間を大幅に短縮する可能性のあるモデルで入力機能を正規化するために広く使用されている手法です。これは近年の機械学習において、より一般的で有用なアルゴリズムの改善の 1 つであり、Inception v3 を含む幅広いモデルで使用されています。
活性化入力は、まずバッチ平均を減算し、バッチ標準偏差で除算することによって正規化されますが、バッチ正規化はそれだけではありません。逆伝播がある場合にバランスを保つため、2 つのトレーニング可能なパラメータが各層に導入されています。正規化された出力 \({\hat{x}}\) は後続の演算 \({\gamma\hat{x}}+\beta\) で処理されます。\({\gamma}\) と \({\beta}\) は標準偏差と平均の一種ですが、これらはモデル自体によって学習されます。
方程式の完全なセットは論文内にありますが、便宜上ここで繰り返します。
入力: ミニバッチの x の値: \(\Phi=\) { \({x_{1..m}\\} \) } 学習対象のパラメータ: \({\gamma}\),\({\beta}\)
出力: { \({y_i}=BN_{\gamma,\beta}{(x_i)}\) }
\[{\mu_\phi} \leftarrow {\frac{1}{m}}{\sum_{i=1}^m}x_i \qquad \mathsf(mini-batch\ mean)\]
\[{\sigma_\phi}^2 \leftarrow {\frac{1}{m}}{\sum_{i=1}^m} {(x_i - {\mu_\phi})^2} \qquad \mathbf(mini-batch\ variance)\]
\[{\hat{x_i}} \leftarrow {\frac{x_i-{\mu_\phi}}{\sqrt {\sigma^2_\phi}+{\epsilon}}}\qquad \mathbf(normalize)\]
\[{y_i}\leftarrow {\gamma \hat{x_i}} + \beta \equiv BN_{\gamma,\beta}{(x_i)}\qquad \mathbf(scale \ and \ shift)\]
正規化はトレーニング中にはうまくいきましたが、評価時には、決定論的な方法でモデルを動作させたいと考えています。画像の分類結果は入力画像のみに依存し、モデルに供給される画像セットには依存しません。したがって、\({\mu}\) と \({\sigma}^2\) を修正する必要があり、これらは画像母集団統計を表す必要があります。
これを行うには、モデルはミニバッチの平均と分散の移動平均を計算します。
\[{\hat\mu_i} = {\alpha \hat\mu_{t-1}}+{(1-\alpha)\mu_t}\]
\[{\hat\sigma_t}^2 = {\alpha{\hat\sigma^2_{t-1}}} + {(1-\alpha) {\sigma_t}^2}\]
Inception v3 の特定のケースでは、GPU で使用するために（ハイパーパラメータ調整を介して）有効な減衰係数を取得していました。TPU でもこの値を使用したいと思いますが、そのためにはいくつかの調整が必要です。
Batchnorm の移動平均と分散は、損失パスフィルタを使用して計算されます。以下に正準方程式を示します（ここでは \({y_t}\) が移動平均または分散を表します）。
\[{y_t}={\alpha y_{t-1}}+{(1-\alpha)}{x_t} \]
（1）
8x1 GPU（同期）ジョブでは、各レプリカは現在の移動平均を読み取り、更新します。新しい移動変数が最初に現在のレプリカによって書き込まれてから、次のレプリカが読み取ることができるという意味では、更新は順次的です。
8 つのレプリカがある場合、アンサンブル更新の一連の演算は次のとおりです。
\[{y_t}={\alpha y_{t-1}}+{(1-\alpha)}{x_t} \]
\[{y_{t+1}}={\alpha y_{t}}+{(1-\alpha)}{x_{t+1}} \]
\[{y_{t+2}}={\alpha y_{t+1}}+{(1-\alpha)}{x_{t+2}} \]
\[{y_{t+3}}={\alpha y_{t+2}}+{(1-\alpha)}{x_{t+3}} \]
\[{y_{t+4}}={\alpha y_{t+3}}+{(1-\alpha)}{x_{t+4}} \]
\[{y_{t+5}}={\alpha y_{t+4}}+{(1-\alpha)}{x_{t+5}} \]
\[{y_{t+6}}={\alpha y_{t+5}}+{(1-\alpha)}{x_{t+6}} \]
\[{y_{t+7}}={\alpha y_{t+6}}+{(1-\alpha)}{x_{t+7}} \]
この一連の連続した 8 つの更新は、次のように表すことができます。
\[{y_{t+7}}={\alpha^8y_{t-1}}+(1-\alpha){\sum_{k=0}^7} {\alpha^{7-k}}{x_{t+k}}\]
（2）
TPU 上の現在の移動モーメント計算の実装では、各シャードは独立して計算を行い、シャード間の通信は行われません。バッチは各シャードに分散され、各シャードはバッチの総数の 1/8（シャードが 8 つある場合）を処理します。
各シャードは指定通りの作業をして、移動モーメント（つまり平均と分散）を計算しますが、シャード 0 の結果のみがホスト CPU に返されます。したがって、事実上、移動平均 / 分散の更新を行っているレプリカは 1 つのみです。
\[{z_t}={\beta {z_{t-1}}}+{(1-\beta)u_t}\]
（3）
この更新は、それに続くシャードの速度の 1/8 で行われています。
GPU と TPU の更新式を比較するには、それぞれの時間尺度を調整する必要があります。具体的には、GPU 上の一連の 8 つの順次更新を含む一連の演算は、TPU 上の単一の更新と比較する必要があります。以下の図にこれを示します。
時間インデックスを変更した方程式を以下に示します。
\[{y_t}={\alpha^8y_{t-1}}+(1-\alpha){\sum_{k=0}^7} {\alpha^{7-k}}{x_{t-k/8}} \qquad \mathsf(GPU)\]
\[{z_t}={\beta {z_{t-1}}}+{(1-\beta)u_t}\qquad \mathsf(TPU) \]
8 つのミニバッチ（関連するすべてのディメンション間で正規化された）がそれぞれ、GPU 上の 8 つのミニバッチの順次更新内で同様の値を生成すると仮定すると、これらの式を次のようにできます。
\[{y_t}={\alpha^8y_{t-1}}+(1-\alpha){\sum_{k=0}^7} {\alpha^{7-k}}{\hat{x_t}}={\alpha^8y_{t-1}+(1-\alpha^8){\hat{x_t}}} \qquad \mathsf(GPU)\]
\[{z_t}={\beta {z_{t-1}}}+{(1-\beta)u_t}\qquad \mathsf(TPU) \]
したがって、GPU に指定された減衰係数の効果に見合うようにするには、それに応じて TPU の減衰係数を変更する必要があります。具体的には、\({\beta}\)=\({\alpha}^8\) を設定する必要があります。
Inception v3 では、GPU で使用される減衰値は \({\alpha}\)=0.9997 です。これにより、TPU の減衰値が \({\beta}\)=0.9976 に変換されます。
学習率の適応
バッチサイズが大きくなるにつれて、トレーニングはより困難になります。大きなバッチサイズで効率的なトレーニングを可能にするさまざまな手法が提案されています。次のページ（ImageNet Training in Minutes、Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour、Extremely Large Minibatch SGD: Training ResNet-50 on ImageNet in 15 Minutes）などをご覧ください。
上記の手法の 1 つ（緩やかな学習率の上昇）を使用してモデルをトレーニングし、4,096～16,384 のバッチサイズについて 78.1% を超える精度を達成しました。Inception v3 の場合、学習率は通常、初期学習率である約 10% に設定されます。学習率は、指定された（小さい）数の「コールド エポック」ではこの低い値のまま一定で、その後、指定された数の「ウォームアップ エポック」では線形増加を開始します。そして最後に、通常の指数関数的減衰が使用されていた場合の学習率と交差します。以下の図にこれを示します。
これを行うコード スニペットを以下に示します。
initial_learning_rate = FLAGS.learning_rate * FLAGS.train_batch_size / 256
if FLAGS.use_learning_rate_warmup:
warmup_decay = FLAGS.learning_rate_decay**(
(FLAGS.warmup_epochs + FLAGS.cold_epochs) /
FLAGS.learning_rate_decay_epochs)
adj_initial_learning_rate = initial_learning_rate * warmup_decay
final_learning_rate = 0.0001 * initial_learning_rate
train_op = None
if training_active:
batches_per_epoch = _NUM_TRAIN_IMAGES / FLAGS.train_batch_size
global_step = tf.train.get_or_create_global_step()
current_epoch = tf.cast(
(tf.cast(global_step, tf.float32) / batches_per_epoch), tf.int32)
learning_rate = tf.train.exponential_decay(
learning_rate=initial_learning_rate,
global_step=global_step,
decay_steps=int(FLAGS.learning_rate_decay_epochs * batches_per_epoch),
decay_rate=FLAGS.learning_rate_decay,
staircase=True)
if FLAGS.use_learning_rate_warmup:
wlr = 0.1 * adj_initial_learning_rate
wlr_height = tf.cast(
0.9 * adj_initial_learning_rate /
(FLAGS.warmup_epochs + FLAGS.learning_rate_decay_epochs - 1),
tf.float32)
epoch_offset = tf.cast(FLAGS.cold_epochs - 1, tf.int32)
exp_decay_start = (FLAGS.warmup_epochs + FLAGS.cold_epochs +
FLAGS.learning_rate_decay_epochs)
lin_inc_lr = tf.add(
wlr, tf.multiply(
tf.cast(tf.subtract(current_epoch, epoch_offset), tf.float32),
wlr_height))
learning_rate = tf.where(
tf.greater_equal(current_epoch, FLAGS.cold_epochs),
(tf.where(tf.greater_equal(current_epoch, exp_decay_start),
learning_rate, lin_inc_lr)),
wlr)
# Set a minimum boundary for the learning rate.
learning_rate = tf.maximum(
learning_rate, final_learning_rate, name='learning_rate')
