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

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

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())
 

以上






AI導入支援 #2 ウェビナー

スモールスタートを可能としたAI導入支援   Vol.2
[無料 WEB セミナー] [詳細]
「画像認識 AI PoC スターターパック」の紹介
既に AI 技術を実ビジネスで活用し、成果を上げている日本企業も多く存在しており、競争優位なビジネスを展開しております。
しかしながら AI を導入したくとも PoC (概念実証) だけでも高額な費用がかかり取組めていない企業も少なくないようです。A I導入時には欠かせない PoC を手軽にしかも短期間で認知度を確認可能とするサービの紹介と共に、AI 技術の特性と具体的な導入プロセスに加え運用時のポイントについても解説いたします。
日時:2021年10月13日(水)
会場:WEBセミナー
共催:クラスキャット、日本FLOW(株)
後援:働き方改革推進コンソーシアム
参加費: 無料 (事前登録制)
人工知能開発支援
◆ クラスキャットは 人工知能研究開発支援 サービスを提供しています :
  • テクニカルコンサルティングサービス
  • 実証実験 (プロトタイプ構築)
  • アプリケーションへの実装
  • 人工知能研修サービス
◆ お問合せ先 ◆
(株)クラスキャット
セールス・インフォメーション
E-Mail:sales-info@classcat.com