Sonnet 2.0 : Tutorials : snt.distribute で分散訓練 (CIFAR-10)

Posted on by Sales Information

Sonnet 2.0 : Tutorials : snt.distribute で分散訓練 (翻訳/解説)

翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 06/10/2020

* 本ページは、Sonnet の以下のドキュメントを翻訳した上で適宜、補足説明したものです：

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

 

snt.distribute で分散訓練

イントロダクション

このチュートリアルは Sonnet 2 “Hello, world!” サンプル (MLP on MNIST) を既に完了していることを仮定しています。

このチュートリアルでは、より大きなモデルとより大きなデータセットで物事をスケールアップしていきます、そして計算をマルチデバイスに渡り分散していきます。

 

序

import sys
assert sys.version_info >= (3, 6), "Sonnet 2 requires Python >=3.6"

!pip install dm-sonnet tqdm

import sonnet as snt
import tensorflow as tf
import tensorflow_datasets as tfds

print("TensorFlow version: {}".format(tf.__version__))
print("    Sonnet version: {}".format(snt.__version__))

最後に利用可能な GPU を素早く見ましょう :

!grep Model: /proc/driver/nvidia/gpus/*/information | awk '{$1="";print$0}'

 

分散ストラテジー

幾つかのデバイスに渡り計算を分散するためのストラテジーが必要です。Google Colab は単一 GPU を提供するだけですのでそれを 4 つの仮想 GPU に分割します :

physical_gpus = tf.config.experimental.list_physical_devices("GPU")
physical_gpus

[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]

tf.config.experimental.set_virtual_device_configuration(
    physical_gpus[0],
    [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2000)] * 4
)

gpus = tf.config.experimental.list_logical_devices("GPU")
gpus

[LogicalDevice(name='/job:localhost/replica:0/task:0/device:GPU:0', device_type='GPU'),
 LogicalDevice(name='/job:localhost/replica:0/task:0/device:GPU:1', device_type='GPU'),
 LogicalDevice(name='/job:localhost/replica:0/task:0/device:GPU:2', device_type='GPU'),
 LogicalDevice(name='/job:localhost/replica:0/task:0/device:GPU:3', device_type='GPU')]

Sonnet optimizer を使用するとき、snt.distribute からの Replicator か TpuReplicator を利用しなければあんりません、あるいは tf.distribute.OneDeviceStrategy を利用できます。Replicator は MirroredStrategy と等値でそして TpuReplicator は TPUStrategy と等値です。

strategy = snt.distribute.Replicator(
    ["/device:GPU:{}".format(i) for i in range(4)],
    tf.distribute.ReductionToOneDevice("GPU:0"))

 

データセット

基本的には MNIST サンプルと同じですが、今回は CIFAR-10 を使用しています。CIFAR-10 は 10 の異なるクラス (飛行機、自動車、鳥、猫、鹿、犬、蛙、馬、船そしてトラック) にある 32×32 ピクセルカラー画像を含みます。

# NOTE: This is the batch size across all GPUs.
batch_size = 100 * 4

def process_batch(images, labels):
  images = tf.cast(images, dtype=tf.float32)
  images = ((images / 255.) - .5) * 2.
  return images, labels

def cifar10(split):
  dataset = tfds.load("cifar10", split=split, as_supervised=True)
  dataset = dataset.map(process_batch)
  dataset = dataset.batch(batch_size)
  dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
  dataset = dataset.cache()
  return dataset

cifar10_train = cifar10("train").shuffle(10)
cifar10_test = cifar10("test")

 

モデル & Optimizer

都合良く、snt.nets にこのデータセットのために特に設計された事前ビルドされたモデルがあります。

作成された任意の変数が正しく分散されることを確実にするために、モデルと optimizer は strategy スコープ内で構築しなければなりません。代わりに、tf.distribute.experimental_set_strategy を使用してプログラム全体のためのスコープに入ることもでできるでしょう。

learning_rate = 0.1

with strategy.scope():
  model = snt.nets.Cifar10ConvNet()
  optimizer = snt.optimizers.Momentum(learning_rate, 0.9)

 

モデルを訓練する

Sonnet optimizer はできる限り綺麗でそして単純であるように設計されています。それらは分散実行を扱うためのどのようなコードも含みません。従ってそれはコードの 2, 3 の追加行を必要とします。

異なるデバイス上で計算された勾配を集めなければなりません。これは ReplicaContext.all_reduce を使用して成されます。

Replicator / TpuReplicator を使用するとき values が総てのレプリカで同一で在り続けることを確かなものにすることはユーザの責任であることに注意してください。

def step(images, labels):
  """Performs a single training step, returning the cross-entropy loss."""
  with tf.GradientTape() as tape:
    logits = model(images, is_training=True)["logits"]
    loss = tf.reduce_mean(
        tf.nn.sparse_softmax_cross_entropy_with_logits(labels=labels,
                                                       logits=logits))

  grads = tape.gradient(loss, model.trainable_variables)

  # Aggregate the gradients from the full batch.
  replica_ctx = tf.distribute.get_replica_context()
  grads = replica_ctx.all_reduce("mean", grads)

  optimizer.apply(grads, model.trainable_variables)
  return loss

@tf.function
def train_step(images, labels):
  per_replica_loss = strategy.run(step, args=(images, labels))
  return strategy.reduce("sum", per_replica_loss, axis=None)

def train_epoch(dataset):
  """Performs one epoch of training, returning the mean cross-entropy loss."""
  total_loss = 0.0
  num_batches = 0

  # Loop over the entire training set.
  for images, labels in dataset:
    total_loss += train_step(images, labels).numpy()
    num_batches += 1

  return total_loss / num_batches

cifar10_train_dist = strategy.experimental_distribute_dataset(cifar10_train)

for epoch in range(20):
  print("Training epoch", epoch, "...", end=" ")
  print("loss :=", train_epoch(cifar10_train_dist))

 

モデルを評価する

バッチ次元に渡り削減するために strategy.reduce による axis パラメータの使用方法に注意してください。

num_cifar10_test_examples = 10000

def is_predicted(images, labels):
  logits = model(images, is_training=False)["logits"]
  # The reduction over the batch happens in `strategy.reduce`, below.
  return tf.cast(tf.equal(labels, tf.argmax(logits, axis=1)), dtype=tf.int32)

cifar10_test_dist = strategy.experimental_distribute_dataset(cifar10_test)

@tf.function
def evaluate():
  """Returns the top-1 accuracy over the entire test set."""
  total_correct = 0

  for images, labels in cifar10_test_dist:
    per_replica_correct = strategy.run(is_predicted, args=(images, labels))
    total_correct += strategy.reduce("sum", per_replica_correct, axis=0)

  return tf.cast(total_correct, tf.float32) / num_cifar10_test_examples

print("Testing...", end=" ")
print("top-1 accuracy =", evaluate().numpy())
 

以上