Inception v3 詳細ガイド

このドキュメントでは、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』という論文をベースにしています。

このモデル自体は、畳み込み層、平均プーリング層、最大プーリング層、連結層、ドロップアウト層、全結合層など、対称と非対称の構成要素で構成されています。バッチ正規化はモデル全体で広く使用され、活性化入力に適用されます。損失は Softmax を使用して計算されます。

次のスクリーンショットは、モデルの概要図です。

イメージ

Estimator API

Inception v3 の TPU バージョンは、開発を容易にするために設計された API である TPUEstimator を使用して作成されているため、基盤となるハードウェアの詳細ではなく、モデルそのものに集中できます。この API は、チェックポイントの保存や復元などの一般的な機能を自動化しながら、バックグラウンドで TPU でモデルを実行するために必要な低水準の重作業のほとんどを行います。

Estimator API は、コードのモデル部分と入力部分を分けます。 モデル定義と入力パイプラインに対応する 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() で、次のコードに示すように、トレーニングと評価に使用されます。

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 データセット

モデルを使用して画像を認識する前に、多数のラベル付き画像を使用してモデルをトレーニングする必要があります。ImageNet は、一般的に使用されるデータセットです。

ImageNet には、1,000 万を超える、ラベル付き画像の URL があります。100 万枚の画像には、ラベル付きオブジェクトのより正確な位置を指定する境界ボックスもあります。

このモデルの場合、ImageNet データセットは 1,331,167 枚の画像で構成され、1,281,167 枚の画像を含むトレーニング用データセットと 50,000 枚の画像を含む評価用のデータセットに分割されています。

トレーニング用と評価用のデータセットは意図的に分割されています。モデルのトレーニングにはトレーニング用データセットの画像のみが使用され、モデルの精度の評価には評価用データセットの画像のみが使用されます。

このモデルは、画像を TFRecord として保存することを想定しています。画像を未加工の JPEG ファイルから TFRecord に変換する方法については、download_and_preprocess_imagenet.sh をご覧ください。

入力パイプライン

各 Cloud TPU デバイスには 8 つのコアがあり、ホスト(CPU)に接続されています。 より大きいスライスには複数のホストがあります。他のより大規模な構成は、複数のホストと通信します。たとえば、v2-256 は 16 のホストと通信します。

ホストは、ファイル システムまたはローカルメモリからデータを取得し、必要なデータ前処理を行い、前処理されたデータを TPU のコアに転送します。 ホストが個別に行うデータ処理の 3 つのフェーズ 1)保存、2)前処理、3)転送について考えます。 概要を以下の図に示します。

画像

良好なパフォーマンスを得るには、システムのバランスを取る必要があります。ホスト CPU が 3 つのデータ処理フェーズを完了するのに TPU よりも時間がかかる場合、実行はホストバウンドになります次の図に、両方のケースを示します。

画像

Inception v3 の現在の実装は、入力バウンドです。画像はファイル システムから取得され、デコードされてから前処理されます。中程度のものから複雑なものまで、さまざまなタイプの前処理ステージを使用できます。最も複雑な前処理ステージを使用すると、トレーニング パイプラインは前処理バウンドになります。モデルを TPU バウンドにする、適度に複雑な前処理ステージを使用して、78.1% を超える精度を達成できます。

このモデルは tf.data.Dataset を使用して入力パイプラインの処理を処理します。入力パイプラインを最適化する方法の詳細については、データセットのパフォーマンス ガイドをご覧ください。

関数を定義して Estimator API に渡すことはできますが、クラス InputPipeline は必要なすべての機能をカプセル化します。

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 の主な要素を示しています。

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.pyinception_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 実行時間バウンドにすることができます。また、1,024 ~ 16,384 のバッチサイズで Inception v3 モデルをトレーニングし、78.1% を超える精度を達成しました。

オプティマイザー

現在のモデルでは、SGD、Momentum、RMSProp という 3 種類のオプティマイザーを紹介しています。

Stochastic gradient descent (SGD) は最も単純な更新であり、重みは負の勾配方向に移動します。その単純さにもかかわらず、一部のモデルでは良好な結果が得られます。更新のダイナミクスは以下のように表すことができます。

$$w_{k+1}=w_k-\alpha∇f(w_k)$$

Momentum は一般的なオプティマイザーであり、SGD よりも早く収束することがあります このオプティマイザーは、SGD と同様に重みを更新しますが、前の更新の方向にもコンポーネントを追加します。次の式は、Momentum オプティマイザーによって実行される更新を示しています。

$$z_{k+1}=\beta z_k+∇f(w_k)$$
$$w_{k+1}=w_k-\alpha z_{k+1}$$

これは、以下のように表すことができます。

$$w_{k+1}=w_k-\alpha ∇f(w_k)+\beta \left(w_k-w_{k-1}\right)$$

最後の項は、前の更新の方向のコンポーネントです。

画像

momentum \({\beta}\) には 0.9 の値を使用します。

RMSprop は、講義で Geoff Hinton が最初に提唱した一般的なオプティマイザーです。 次の式は、オプティマイザーの仕組みを表します。

$$g_{k+1}^{-2} = \alpha g_{k}^{-2} + (1-\alpha) g_{k}^2$$ $$w_{k+1}=\beta w_k + \frac{\eta}{\sqrt {g_{k+1^{\mathbf{+{\epsilon}}}}^{-2}}} ∇f(w_k)$$

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)

指数移動平均

トレーニング中、トレーニング可能なパラメータは、オプティマイザーの更新ルールに従ってバックプロパゲーション中に更新されます。これらのルールを記述する式は前のセクションで説明しましたが、便宜上ここでも繰り返します。

$${\theta_{k+1}} = {\theta_k}-{\alpha ∇f(\theta_k)} \qquad(SGD)$$
$${\theta_{k+1}}={\theta_k}-{\alpha z_{k+1}} \qquad(Momentum)$$
$${\theta_{k+1}}= {\beta \theta_k}+\frac{\eta}{\sqrt {g_{k+1^{\mathbf+{\epsilon}}}}^{-2}} ∇f(\theta_k) \qquad(RMSprop)$$

指数移動平均(指数平滑とも呼ばれます)は、更新された重みに適用されるオプションの後処理ステップであり、パフォーマンスの著しい向上につながることがあります。TensorFlow は関数 tf.train.ExponentialMovingAverage を提供し、以下の数式を使用して、重み \({\theta}\) の指数移動平均 \({\hat{\theta}}\) を計算します。

$${\hat{\theta_t}}={\alpha {\hat{\theta}{_{t-1}}}}+{(1-\alpha)}{\theta_t}$$

\({\alpha}\) は減衰係数(1.0 に近い)です。Inception v3 モデルでは、\({\alpha}\) は 0.995 に設定されています。

この計算は無限インパルスレスポンス(IIR)フィルタですが、減衰係数は、次の図に示すように、ほとんどのエネルギー(関連するサンプル)が存在する有効なウィンドウを確立します。

画像

次のように、フィルタ式を書き換えることができます。

$${\hat{\theta}_{t+T+1}}={\alpha(1-\alpha)}({\theta_{t+T}}+{\alpha \theta_{t+T-1}}+...+{\alpha^{t+T}}{\theta_0})$$

ここでは \({\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 でもこの値を使用したいと思いますが、そのためにはいくつかの調整が必要です。

バッチ正規化の移動平均と分散は、次の式に示すように損失パスフィルタを使用して計算されます(ここでは \({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 MinutesAccurate, Large Minibatch SGD: Training ImageNet in 1 HourExtremely 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')