Catat tanggalnya! Google I / O mengembalikan 18-20 Mei Daftar sekarang
Halaman ini diterjemahkan oleh Cloud Translation API.
Switch to English

Gunakan TPU

Lihat di TensorFlow.org Jalankan di Google Colab Lihat sumber di GitHub Unduh buku catatan

Dukungan eksperimental untuk Cloud TPU saat ini tersedia untuk Keras dan Google Colab. Sebelum Anda menjalankan notebook Colab ini, pastikan akselerator hardware Anda adalah TPU dengan memeriksa setelan notebook Anda: Runtime> Ubah jenis runtime> Akselerator hardware> TPU.

Mempersiapkan

import tensorflow as tf

import os
import tensorflow_datasets as tfds

Inisialisasi TPU

TPU biasanya pada pekerja Cloud TPU yang berbeda dari proses lokal yang menjalankan program python pengguna. Oleh karena itu, beberapa pekerjaan inisialisasi perlu dilakukan untuk terhubung ke cluster jarak jauh dan menginisialisasi TPU. Perhatikan bahwa argumen tpu ke TPUClusterResolver adalah alamat khusus hanya untuk Colab. Jika Anda menjalankan Google Compute Engine (GCE), Anda harus meneruskan nama CloudTPU Anda.

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

INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.74:8470
INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.74:8470
INFO:tensorflow:Clearing out eager caches
INFO:tensorflow:Clearing out eager caches
INFO:tensorflow:Finished initializing TPU system.
INFO:tensorflow:Finished initializing TPU system.
All devices:  [LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:7', 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:5', 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:3', device_type='TPU'), 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')]

Penempatan perangkat manual

Setelah TPU diinisialisasi, Anda dapat menggunakan penempatan perangkat manual untuk menempatkan komputasi pada satu perangkat TPU.

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)

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

Strategi distribusi

Sering kali pengguna ingin menjalankan model pada beberapa TPU secara paralel data. Strategi distribusi adalah abstraksi yang dapat digunakan untuk menjalankan model di CPU, GPU, atau TPU. Cukup tukar strategi distribusi dan model akan berjalan di perangkat yang diberikan. Lihat panduan strategi distribusi untuk informasi lebih lanjut.

Pertama, membuat objek TPUStrategy .

strategy = tf.distribute.TPUStrategy(resolver)

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)

Untuk mereplikasi komputasi agar dapat berjalan di semua inti TPU, Anda cukup meneruskannya ke strategy.run API. Di bawah ini adalah contoh bahwa semua core akan mendapatkan input yang sama (a, b) , dan melakukan matmul pada setiap core secara independen. Keluarannya adalah nilai dari semua replika.

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

z = strategy.run(matmul_fn, args=(a, b))
print(z)

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

Klasifikasi pada TPU

Karena kita telah mempelajari konsep dasar, sekarang saatnya untuk melihat contoh yang lebih konkret. Panduan ini menunjukkan cara menggunakan strategi distribusi tf.distribute.TPUStrategy untuk menjalankan Cloud TPU dan melatih model Keras.

Tentukan model Keras

Di bawah ini adalah definisi model MNIST yang menggunakan Keras, tidak berubah dari apa yang akan Anda gunakan pada CPU atau GPU. Perhatikan bahwa pembuatan model Keras harus berada di dalam strategy.scope , sehingga variabel dapat dibuat di setiap perangkat TPU. Bagian lain dari kode tidak perlu berada di dalam cakupan strategi.

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

Masukkan set data

Penggunaan yang efisien daritf.data.Dataset API sangat penting saat menggunakan Cloud TPU, karena Cloud TPU tidak dapat digunakan kecuali Anda dapat memberi mereka data dengan cukup cepat. Lihat Panduan Kinerja Saluran Masukan untuk detail tentang kinerja set data.

Untuk semua kecuali eksperimen paling sederhana (menggunakan tf.data.Dataset.from_tensor_slices atau data dalam grafik lainnya), Anda perlu menyimpan semua file data yang dibaca oleh Set Data di bucket Google Cloud Storage (GCS).

Untuk sebagian besar kasus penggunaan, disarankan untuk mengonversi data Anda ke dalam format TFRecord dan menggunakan tf.data.TFRecordDataset untuk membacanya. Lihat TFRecord dan tf. Contoh tutorial untuk detil bagaimana melakukan ini. Namun, ini bukan persyaratan yang sulit dan Anda dapat menggunakan pembaca FixedLengthRecordDataset data lain ( FixedLengthRecordDataset atau TextLineDataset ) jika mau.

tf.data.Dataset.cache data kecil dapat dimuat seluruhnya ke dalam memori menggunakan tf.data.Dataset.cache .

Terlepas dari format data yang digunakan, sangat disarankan agar Anda menggunakan file besar, dengan urutan 100MB. Hal ini sangat penting dalam pengaturan jaringan ini karena overhead pembukaan file jauh lebih tinggi.

Di sini Anda harus menggunakan modul tensorflow_datasets untuk mendapatkan salinan data pelatihan MNIST. Perhatikan bahwa try_gcs ditentukan untuk menggunakan salinan yang tersedia di bucket GCS publik. Jika Anda tidak menentukannya, TPU tidak akan dapat mengakses data yang diunduh.

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)

  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 to have a
  # infinite dataset for training is to avoid the potential last partial batch
  # in each epoch, so users 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

Latih model menggunakan API tingkat tinggi Keras

Anda dapat melatih model hanya dengan Keras fit / compile API. Tidak ada yang spesifik untuk TPU, Anda akan menulis kode yang sama di bawah ini jika Anda memiliki GPU mutliple dan menggunakan MirroredStrategy daripada TPUStrategy . Untuk mempelajari lebih lanjut, lihat tutorial pelatihan Terdistribusi dengan Keras .

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)

Epoch 1/5
300/300 [==============================] - 19s 38ms/step - loss: 0.3196 - sparse_categorical_accuracy: 0.8996 - val_loss: 0.0467 - val_sparse_categorical_accuracy: 0.9855
Epoch 2/5
300/300 [==============================] - 7s 25ms/step - loss: 0.0381 - sparse_categorical_accuracy: 0.9878 - val_loss: 0.0379 - val_sparse_categorical_accuracy: 0.9885
Epoch 3/5
300/300 [==============================] - 8s 26ms/step - loss: 0.0213 - sparse_categorical_accuracy: 0.9933 - val_loss: 0.0366 - val_sparse_categorical_accuracy: 0.9892
Epoch 4/5
300/300 [==============================] - 8s 25ms/step - loss: 0.0130 - sparse_categorical_accuracy: 0.9961 - val_loss: 0.0327 - val_sparse_categorical_accuracy: 0.9902
Epoch 5/5
300/300 [==============================] - 8s 25ms/step - loss: 0.0083 - sparse_categorical_accuracy: 0.9975 - val_loss: 0.0465 - val_sparse_categorical_accuracy: 0.9885
<tensorflow.python.keras.callbacks.History at 0x7fc0906c90b8>

Untuk mengurangi overhead python, dan memaksimalkan kinerja TPU Anda, coba argumen eksperimental experimental_steps_per_execution ke Model.compile . Di sini ia meningkatkan throughput sekitar 50%:

with strategy.scope():
  model = create_model()
  model.compile(optimizer='adam',
                # Anything between 2 and `steps_per_epoch` could help here.
                experimental_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)

WARNING:tensorflow:The argument `steps_per_execution` is no longer experimental. Pass `steps_per_execution` instead of `experimental_steps_per_execution`.
WARNING:tensorflow:The argument `steps_per_execution` is no longer experimental. Pass `steps_per_execution` instead of `experimental_steps_per_execution`.
Epoch 1/5
300/300 [==============================] - 14s 46ms/step - loss: 0.2399 - sparse_categorical_accuracy: 0.9279 - val_loss: 0.0510 - val_sparse_categorical_accuracy: 0.9838
Epoch 2/5
300/300 [==============================] - 5s 15ms/step - loss: 0.0406 - sparse_categorical_accuracy: 0.9876 - val_loss: 0.0409 - val_sparse_categorical_accuracy: 0.9882
Epoch 3/5
300/300 [==============================] - 5s 15ms/step - loss: 0.0203 - sparse_categorical_accuracy: 0.9936 - val_loss: 0.0394 - val_sparse_categorical_accuracy: 0.9879
Epoch 4/5
300/300 [==============================] - 5s 15ms/step - loss: 0.0132 - sparse_categorical_accuracy: 0.9959 - val_loss: 0.0410 - val_sparse_categorical_accuracy: 0.9879
Epoch 5/5
300/300 [==============================] - 5s 15ms/step - loss: 0.0140 - sparse_categorical_accuracy: 0.9953 - val_loss: 0.0508 - val_sparse_categorical_accuracy: 0.9863
<tensorflow.python.keras.callbacks.History at 0x7fbc743646d8>

Latih model menggunakan loop pelatihan kustom.

Anda juga dapat membuat dan melatih model Anda menggunakan tf.function dan tf.distribute API secara langsung. strategy.experimental_distribute_datasets_from_function API digunakan untuk mendistribusikan dataset yang diberi fungsi dataset. Perhatikan bahwa ukuran batch yang diteruskan ke set data akan menjadi ukuran batch per replika, bukan ukuran batch global dalam kasus ini. Untuk mempelajari lebih lanjut, lihat tutorial Custom training dengan tf.distribute.Strategy tutorial.

Pertama, buat model, dataset dan tf.functions.

# Create the model, optimizer and metrics inside 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),))

WARNING:tensorflow:From <ipython-input-1-2f075cabff81>: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-2f075cabff81>: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

Kemudian jalankan loop pelatihan.

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

Epoch: 0/5
Current step: 300, training loss: 0.1237, accuracy: 96.17%
Epoch: 1/5
Current step: 600, training loss: 0.0333, accuracy: 98.98%
Epoch: 2/5
Current step: 900, training loss: 0.0185, accuracy: 99.37%
Epoch: 3/5
Current step: 1200, training loss: 0.0135, accuracy: 99.54%
Epoch: 4/5
Current step: 1500, training loss: 0.0092, accuracy: 99.7%

Meningkatkan kinerja dengan beberapa langkah dalam tf.function

Performa dapat ditingkatkan dengan menjalankan beberapa langkah dalam fungsi tf.function . Ini dicapai dengan tf.range panggilan strategy.run dengan tf.range di dalam tf.function , AutoGraph akan mengubahnya menjadi tf.while_loop pada pekerja TPU.

Meskipun dengan kinerja yang lebih baik, ada tf.function dibandingkan dengan satu langkah di dalam tf.function . Menjalankan beberapa langkah dalam sebuah tf.function kurang fleksibel, Anda tidak dapat menjalankan sesuatu dengan bersemangat atau kode python sewenang-wenang dalam langkah-langkah tersebut.

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

Current step: 1800, training loss: 0.009, accuracy: 99.71%

Langkah selanjutnya