TPU'ları kullanın

TensorFlow.org'da görüntüleyin Google Colab'da çalıştırın Kaynağı GitHub'da görüntüleyin Not defterini indir

Bu Colab not defterini çalıştırmadan önce, not defteri ayarlarınızı kontrol ederek donanım hızlandırıcınızın bir TPU olduğundan emin olun: Çalışma zamanı > Çalışma zamanı türünü değiştir > Donanım hızlandırıcı > TPU .

Kurmak

import tensorflow as tf

import os
import tensorflow_datasets as tfds
tutucu1 l10n-yer
/tmpfs/src/tf_docs_env/lib/python3.6/site-packages/requests/__init__.py:104: RequestsDependencyWarning: urllib3 (1.26.8) or chardet (2.3.0)/charset_normalizer (2.0.11) doesn't match a supported version!
  RequestsDependencyWarning)

TPU başlatma

TPU'lar, genellikle, kullanıcının Python programını çalıştıran yerel süreçten farklı olan Cloud TPU çalışanlarıdır. Bu nedenle, uzak kümeye bağlanmak ve TPU'ları başlatmak için bazı başlatma çalışmaları yapmanız gerekir. tf.distribute.cluster_resolver.TPUClusterResolver için tpu bağımsız değişkeninin yalnızca Colab için özel bir adres olduğunu unutmayın. Kodunuzu Google Compute Engine'de (GCE) çalıştırıyorsanız, bunun yerine Cloud TPU'nuzun adını iletmelisiniz.

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
# This is the TPU initialization code that has to be at the beginning.
tf.tpu.experimental.initialize_tpu_system(resolver)
print("All devices: ", tf.config.list_logical_devices('TPU'))
tutucu3 l10n-yer
INFO:tensorflow:Clearing out eager caches
INFO:tensorflow:Clearing out eager caches
INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.10:8470
INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.10:8470
INFO:tensorflow:Finished initializing TPU system.
INFO:tensorflow:Finished initializing TPU system.
All devices:  [LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:0', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:1', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:2', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:3', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:4', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:5', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:6', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:7', device_type='TPU')]

Manuel cihaz yerleştirme

TPU başlatıldıktan sonra, hesaplamayı tek bir TPU aygıtına yerleştirmek için manuel aygıt yerleştirmeyi kullanabilirsiniz:

a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])

with tf.device('/TPU:0'):
  c = tf.matmul(a, b)

print("c device: ", c.device)
print(c)
tutucu5 l10n-yer
c device:  /job:worker/replica:0/task:0/device:TPU:0
tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)

Dağıtım stratejileri

Genellikle modelinizi veri paralel bir şekilde birden çok TPU üzerinde çalıştırırsınız. Modelinizi birden çok TPU'ya (veya diğer hızlandırıcılara) dağıtmak için TensorFlow çeşitli dağıtım stratejileri sunar. Dağıtım stratejinizi değiştirebilirsiniz ve model herhangi bir (TPU) cihazda çalışacaktır. Daha fazla bilgi için dağıtım stratejisi kılavuzuna bakın.

Bunu göstermek için bir tf.distribute.TPUStrategy nesnesi oluşturun:

strategy = tf.distribute.TPUStrategy(resolver)
tutucu7 l10n-yer
INFO:tensorflow:Found TPU system:
INFO:tensorflow:Found TPU system:
INFO:tensorflow:*** Num TPU Cores: 8
INFO:tensorflow:*** Num TPU Cores: 8
INFO:tensorflow:*** Num TPU Workers: 1
INFO:tensorflow:*** Num TPU Workers: 1
INFO:tensorflow:*** Num TPU Cores Per Worker: 8
INFO:tensorflow:*** Num TPU Cores Per Worker: 8
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)

Tüm TPU çekirdeklerinde çalışabilmesi için bir hesaplamayı çoğaltmak için, onu strategy.run API'sine geçirebilirsiniz. Aşağıda, aynı girdileri (a, b) alan ve her bir çekirdek üzerinde bağımsız olarak matris çarpımı gerçekleştiren tüm çekirdekleri gösteren bir örnek verilmiştir. Çıktılar, tüm kopyalardan alınan değerler olacaktır.

@tf.function
def matmul_fn(x, y):
  z = tf.matmul(x, y)
  return z

z = strategy.run(matmul_fn, args=(a, b))
print(z)
tutucu9 l10n-yer
PerReplica:{
  0: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  1: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  2: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  3: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  4: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  5: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  6: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  7: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)
}

TPU'larda sınıflandırma

Temel kavramları ele aldıktan sonra, daha somut bir örnek düşünün. Bu bölüm, bir Keras modelini bir Cloud TPU üzerinde eğitmek için dağıtım stratejisinin tf.distribute.TPUStrategy ) nasıl kullanılacağını gösterir.

Bir Keras modeli tanımlayın

Keras kullanarak MNIST veri kümesinde görüntü sınıflandırması için Sequential Keras modelinin tanımıyla başlayın. CPU'lar veya GPU'lar üzerinde eğitim alıyor olsaydınız, kullanacağınızdan farklı değil. Değişkenlerin her TPU cihazında oluşturulabilmesi için Keras modeli oluşturmanın strategy.scope içinde olması gerektiğini unutmayın. Kodun diğer bölümlerinin strateji kapsamı içinde olması gerekli değildir.

def create_model():
  return tf.keras.Sequential(
      [tf.keras.layers.Conv2D(256, 3, activation='relu', input_shape=(28, 28, 1)),
       tf.keras.layers.Conv2D(256, 3, activation='relu'),
       tf.keras.layers.Flatten(),
       tf.keras.layers.Dense(256, activation='relu'),
       tf.keras.layers.Dense(128, activation='relu'),
       tf.keras.layers.Dense(10)])

Veri kümesini yükleyin

Bulut TPU'ları kullanırken tf.data.Dataset API'sinin verimli kullanımı, verileri yeterince hızlı bir şekilde besleyemediğiniz sürece Cloud TPU'ları kullanmak imkansız olduğundan kritik önem taşır. Giriş işlem hattı performans kılavuzunda veri kümesi performansı hakkında daha fazla bilgi edinebilirsiniz.

En basit deneyler dışındaki tüm deneyler için ( tf.data.Dataset.from_tensor_slices veya diğer grafik içi verileri kullanarak), Veri Kümesi tarafından Google Bulut Depolama (GCS) paketlerinde okunan tüm veri dosyalarını depolamanız gerekir.

Çoğu kullanım durumu için, verilerinizi TFRecord biçimine dönüştürmeniz ve okumak için bir tf.data.TFRecordDataset kullanmanız önerilir. Bunun nasıl yapılacağına ilişkin ayrıntılar için TFRecord ve tf.Example eğitimine bakın. Bu zor bir gereklilik değildir ve tf.data.FixedLengthRecordDataset veya tf.data.TextLineDataset gibi diğer veri kümesi okuyucularını kullanabilirsiniz.

tf.data.Dataset.cache kullanarak tüm küçük veri kümelerini belleğe yükleyebilirsiniz.

Kullanılan veri formatı ne olursa olsun, 100MB mertebesinde büyük dosyaları kullanmanız şiddetle tavsiye edilir. Bir dosyayı açmanın ek yükü önemli ölçüde daha yüksek olduğundan, bu ağ bağlantılı ortamda özellikle önemlidir.

Aşağıdaki kodda gösterildiği gibi, MNIST eğitim ve test verilerinin bir kopyasını almak için tensorflow_datasets modülünü kullanmalısınız. try_gcs , genel bir GCS paketinde bulunan bir kopyayı kullanmak üzere belirtildiğini unutmayın. Bunu belirtmezseniz TPU indirilen verilere erişemez.

def get_dataset(batch_size, is_training=True):
  split = 'train' if is_training else 'test'
  dataset, info = tfds.load(name='mnist', split=split, with_info=True,
                            as_supervised=True, try_gcs=True)

  # Normalize the input data.
  def scale(image, label):
    image = tf.cast(image, tf.float32)
    image /= 255.0
    return image, label

  dataset = dataset.map(scale)

  # Only shuffle and repeat the dataset in training. The advantage of having an
  # infinite dataset for training is to avoid the potential last partial batch
  # in each epoch, so that you don't need to think about scaling the gradients
  # based on the actual batch size.
  if is_training:
    dataset = dataset.shuffle(10000)
    dataset = dataset.repeat()

  dataset = dataset.batch(batch_size)

  return dataset

Modeli Keras üst düzey API'lerini kullanarak eğitin

Modelinizi Keras fit ve compile API'leri ile eğitebilirsiniz. Bu adımda TPU'ya özgü hiçbir şey yoktur; kodu, TPUStrategy yerine birden fazla GPU ve MirroredStrategy kullanıyormuşsunuz gibi yazarsınız. Keras öğreticisiyle Dağıtılmış eğitimde daha fazla bilgi edinebilirsiniz.

with strategy.scope():
  model = create_model()
  model.compile(optimizer='adam',
                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=['sparse_categorical_accuracy'])

batch_size = 200
steps_per_epoch = 60000 // batch_size
validation_steps = 10000 // batch_size

train_dataset = get_dataset(batch_size, is_training=True)
test_dataset = get_dataset(batch_size, is_training=False)

model.fit(train_dataset,
          epochs=5,
          steps_per_epoch=steps_per_epoch,
          validation_data=test_dataset, 
          validation_steps=validation_steps)
tutucu13 l10n-yer
Epoch 1/5
300/300 [==============================] - 18s 32ms/step - loss: 0.1433 - sparse_categorical_accuracy: 0.9564 - val_loss: 0.0452 - val_sparse_categorical_accuracy: 0.9859
Epoch 2/5
300/300 [==============================] - 6s 21ms/step - loss: 0.0335 - sparse_categorical_accuracy: 0.9898 - val_loss: 0.0318 - val_sparse_categorical_accuracy: 0.9899
Epoch 3/5
300/300 [==============================] - 6s 21ms/step - loss: 0.0199 - sparse_categorical_accuracy: 0.9935 - val_loss: 0.0397 - val_sparse_categorical_accuracy: 0.9866
Epoch 4/5
300/300 [==============================] - 6s 21ms/step - loss: 0.0109 - sparse_categorical_accuracy: 0.9964 - val_loss: 0.0436 - val_sparse_categorical_accuracy: 0.9892
Epoch 5/5
300/300 [==============================] - 6s 21ms/step - loss: 0.0103 - sparse_categorical_accuracy: 0.9963 - val_loss: 0.0481 - val_sparse_categorical_accuracy: 0.9881
<keras.callbacks.History at 0x7f0d485602e8>

Python ek yükünü azaltmak ve TPU'nuzun performansını en üst düzeye çıkarmak için, step_per_execution argümanını steps_per_execution Model.compile . Bu örnekte, verimi yaklaşık %50 oranında artırır:

with strategy.scope():
  model = create_model()
  model.compile(optimizer='adam',
                # Anything between 2 and `steps_per_epoch` could help here.
                steps_per_execution = 50,
                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=['sparse_categorical_accuracy'])

model.fit(train_dataset,
          epochs=5,
          steps_per_epoch=steps_per_epoch,
          validation_data=test_dataset,
          validation_steps=validation_steps)
tutucu15 l10n-yer
Epoch 1/5
300/300 [==============================] - 12s 41ms/step - loss: 0.1515 - sparse_categorical_accuracy: 0.9537 - val_loss: 0.0416 - val_sparse_categorical_accuracy: 0.9863
Epoch 2/5
300/300 [==============================] - 3s 10ms/step - loss: 0.0366 - sparse_categorical_accuracy: 0.9891 - val_loss: 0.0410 - val_sparse_categorical_accuracy: 0.9875
Epoch 3/5
300/300 [==============================] - 3s 10ms/step - loss: 0.0191 - sparse_categorical_accuracy: 0.9938 - val_loss: 0.0432 - val_sparse_categorical_accuracy: 0.9865
Epoch 4/5
300/300 [==============================] - 3s 10ms/step - loss: 0.0141 - sparse_categorical_accuracy: 0.9951 - val_loss: 0.0447 - val_sparse_categorical_accuracy: 0.9875
Epoch 5/5
300/300 [==============================] - 3s 11ms/step - loss: 0.0093 - sparse_categorical_accuracy: 0.9968 - val_loss: 0.0426 - val_sparse_categorical_accuracy: 0.9884
<keras.callbacks.History at 0x7f0d0463cd68>

Özel bir eğitim döngüsü kullanarak modeli eğitin

Ayrıca, tf.function ve tf.distribute API'lerini doğrudan kullanarak modelinizi oluşturabilir ve eğitebilirsiniz. Bir veri kümesi işlevi verilen veri kümesini dağıtmak için strategy.experimental_distribute_datasets_from_function API'sini kullanabilirsiniz. Aşağıdaki örnekte, veri kümesine iletilen toplu iş boyutunun, genel toplu iş boyutu yerine çoğaltma başına toplu iş boyutu olduğuna dikkat edin. Daha fazla bilgi edinmek için tf.distribute.Strategy öğreticisiyle Özel eğitime göz atın .

İlk önce modeli, veri kümelerini ve tf.işlevlerini oluşturun:

# Create the model, optimizer and metrics inside the strategy scope, so that the
# variables can be mirrored on each device.
with strategy.scope():
  model = create_model()
  optimizer = tf.keras.optimizers.Adam()
  training_loss = tf.keras.metrics.Mean('training_loss', dtype=tf.float32)
  training_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
      'training_accuracy', dtype=tf.float32)

# Calculate per replica batch size, and distribute the datasets on each TPU
# worker.
per_replica_batch_size = batch_size // strategy.num_replicas_in_sync

train_dataset = strategy.experimental_distribute_datasets_from_function(
    lambda _: get_dataset(per_replica_batch_size, is_training=True))

@tf.function
def train_step(iterator):
  """The step function for one training step."""

  def step_fn(inputs):
    """The computation to run on each TPU device."""
    images, labels = inputs
    with tf.GradientTape() as tape:
      logits = model(images, training=True)
      loss = tf.keras.losses.sparse_categorical_crossentropy(
          labels, logits, from_logits=True)
      loss = tf.nn.compute_average_loss(loss, global_batch_size=batch_size)
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))
    training_loss.update_state(loss * strategy.num_replicas_in_sync)
    training_accuracy.update_state(labels, logits)

  strategy.run(step_fn, args=(next(iterator),))
tutucu17 l10n-yer
WARNING:tensorflow:From <ipython-input-1-5625c2a14441>:15: StrategyBase.experimental_distribute_datasets_from_function (from tensorflow.python.distribute.distribute_lib) is deprecated and will be removed in a future version.
Instructions for updating:
rename to distribute_datasets_from_function
WARNING:tensorflow:From <ipython-input-1-5625c2a14441>:15: StrategyBase.experimental_distribute_datasets_from_function (from tensorflow.python.distribute.distribute_lib) is deprecated and will be removed in a future version.
Instructions for updating:
rename to distribute_datasets_from_function

Ardından, eğitim döngüsünü çalıştırın:

steps_per_eval = 10000 // batch_size

train_iterator = iter(train_dataset)
for epoch in range(5):
  print('Epoch: {}/5'.format(epoch))

  for step in range(steps_per_epoch):
    train_step(train_iterator)
  print('Current step: {}, training loss: {}, accuracy: {}%'.format(
      optimizer.iterations.numpy(),
      round(float(training_loss.result()), 4),
      round(float(training_accuracy.result()) * 100, 2)))
  training_loss.reset_states()
  training_accuracy.reset_states()
tutucu19 l10n-yer
Epoch: 0/5
Current step: 300, training loss: 0.1339, accuracy: 95.79%
Epoch: 1/5
Current step: 600, training loss: 0.0333, accuracy: 98.91%
Epoch: 2/5
Current step: 900, training loss: 0.0176, accuracy: 99.43%
Epoch: 3/5
Current step: 1200, training loss: 0.0126, accuracy: 99.61%
Epoch: 4/5
Current step: 1500, training loss: 0.0122, accuracy: 99.61%

tf.function içindeki birden çok adımla performansı iyileştirme

Bir tf.function içinde birden çok adım çalıştırarak performansı artırabilirsiniz. Bu, strategy.run çağrısını tf.range içinde bir tf.function ile sarmalayarak elde edilir ve Otomatik Grafik, TPU çalışanında bunu bir tf.while_loop dönüştürür.

Geliştirilmiş performansa rağmen, bu yöntemle tf.function içinde tek bir adım çalıştırmaya kıyasla bazı ödünleşimler vardır. Bir tf.function içinde birden fazla adım çalıştırmak daha az esnektir; adımları hevesle veya rastgele Python kodunu çalıştıramazsınız.

@tf.function
def train_multiple_steps(iterator, steps):
  """The step function for one training step."""

  def step_fn(inputs):
    """The computation to run on each TPU device."""
    images, labels = inputs
    with tf.GradientTape() as tape:
      logits = model(images, training=True)
      loss = tf.keras.losses.sparse_categorical_crossentropy(
          labels, logits, from_logits=True)
      loss = tf.nn.compute_average_loss(loss, global_batch_size=batch_size)
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))
    training_loss.update_state(loss * strategy.num_replicas_in_sync)
    training_accuracy.update_state(labels, logits)

  for _ in tf.range(steps):
    strategy.run(step_fn, args=(next(iterator),))

# Convert `steps_per_epoch` to `tf.Tensor` so the `tf.function` won't get 
# retraced if the value changes.
train_multiple_steps(train_iterator, tf.convert_to_tensor(steps_per_epoch))

print('Current step: {}, training loss: {}, accuracy: {}%'.format(
      optimizer.iterations.numpy(),
      round(float(training_loss.result()), 4),
      round(float(training_accuracy.result()) * 100, 2)))
tutucu21 l10n-yer
Current step: 1800, training loss: 0.0081, accuracy: 99.74%

Sonraki adımlar