 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='grpc://' + os.environ['COLAB_TPU_ADDR'])
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.2:8470 INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.2: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: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')]
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, TPUStrategy 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:localhost/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:XLA_CPU:0, XLA_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 inti akan mendapatkan masukan yang sama (a, b) , dan melakukan matmul pada setiap inti 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.experimental.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 dari tf.data.Dataset API sangat penting saat menggunakan Cloud TPU, karena Cloud TPU tidak mungkin 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 yang 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 detail tentang 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 di sini yang spesifik untuk TPU, Anda akan menulis kode yang sama di bawah ini jika Anda memiliki beberapa GPU 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 WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/data/ops/multi_device_iterator_ops.py:601: get_next_as_optional (from tensorflow.python.data.ops.iterator_ops) is deprecated and will be removed in a future version. Instructions for updating: Use `tf.data.Iterator.get_next_as_optional()` instead. Warning:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/data/ops/multi_device_iterator_ops.py:601: get_next_as_optional (from tensorflow.python.data.ops.iterator_ops) is deprecated and will be removed in a future version. Instructions for updating: Use `tf.data.Iterator.get_next_as_optional()` instead. 1/300 [..............................] - ETA: 20:25 - loss: 2.3062 - sparse_categorical_accuracy: 0.0500WARNING:tensorflow:Callbacks method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0015s vs `on_train_batch_end` time: 0.0237s). Check your callbacks. Warning:tensorflow:Callbacks method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0015s vs `on_train_batch_end` time: 0.0237s). Check your callbacks. 300/300 [==============================] - ETA: 0s - loss: 0.1450 - sparse_categorical_accuracy: 0.9561WARNING:tensorflow:Callbacks method `on_test_batch_end` is slow compared to the batch time (batch time: 0.0013s vs `on_test_batch_end` time: 0.0097s). Check your callbacks. Warning:tensorflow:Callbacks method `on_test_batch_end` is slow compared to the batch time (batch time: 0.0013s vs `on_test_batch_end` time: 0.0097s). Check your callbacks. 300/300 [==============================] - 15s 50ms/step - loss: 0.1450 - sparse_categorical_accuracy: 0.9561 - val_loss: 0.0456 - val_sparse_categorical_accuracy: 0.9852 Epoch 2/5 300/300 [==============================] - 8s 28ms/step - loss: 0.0357 - sparse_categorical_accuracy: 0.9890 - val_loss: 0.0372 - val_sparse_categorical_accuracy: 0.9884 Epoch 3/5 300/300 [==============================] - 9s 28ms/step - loss: 0.0193 - sparse_categorical_accuracy: 0.9940 - val_loss: 0.0557 - val_sparse_categorical_accuracy: 0.9835 Epoch 4/5 300/300 [==============================] - 9s 29ms/step - loss: 0.0141 - sparse_categorical_accuracy: 0.9954 - val_loss: 0.0405 - val_sparse_categorical_accuracy: 0.9883 Epoch 5/5 300/300 [==============================] - 9s 29ms/step - loss: 0.0092 - sparse_categorical_accuracy: 0.9967 - val_loss: 0.0428 - val_sparse_categorical_accuracy: 0.9887 <tensorflow.python.keras.callbacks.History at 0x7f480c793240>
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)
Epoch 1/5 300/300 [==============================] - 16s 52ms/step - loss: 0.1337 - sparse_categorical_accuracy: 0.9583 - val_loss: 0.0521 - val_sparse_categorical_accuracy: 0.9840 Epoch 2/5 300/300 [==============================] - 5s 16ms/step - loss: 0.0331 - sparse_categorical_accuracy: 0.9898 - val_loss: 0.0360 - val_sparse_categorical_accuracy: 0.9884 Epoch 3/5 300/300 [==============================] - 5s 16ms/step - loss: 0.0188 - sparse_categorical_accuracy: 0.9939 - val_loss: 0.0405 - val_sparse_categorical_accuracy: 0.9887 Epoch 4/5 300/300 [==============================] - 5s 16ms/step - loss: 0.0117 - sparse_categorical_accuracy: 0.9961 - val_loss: 0.0808 - val_sparse_categorical_accuracy: 0.9820 Epoch 5/5 300/300 [==============================] - 5s 16ms/step - loss: 0.0119 - sparse_categorical_accuracy: 0.9962 - val_loss: 0.0488 - val_sparse_categorical_accuracy: 0.9862 <tensorflow.python.keras.callbacks.History at 0x7f47ac791240>
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),))
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.1335, accuracy: 95.8% Epoch: 1/5 Current step: 600, training loss: 0.0344, accuracy: 98.93% Epoch: 2/5 Current step: 900, training loss: 0.0196, accuracy: 99.35% Epoch: 3/5 Current step: 1200, training loss: 0.0119, accuracy: 99.61% Epoch: 4/5 Current step: 1500, training loss: 0.0102, accuracy: 99.65%
Meningkatkan kinerja dengan beberapa langkah dalam tf.function
Performa dapat ditingkatkan dengan menjalankan beberapa langkah dalam fungsi tf.function . Hal 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 sembarang kode python 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.0087, accuracy: 99.72%
Langkah selanjutnya
- Dokumentasi Google Cloud TPU - Siapkan dan jalankan Google Cloud TPU.
- Pelatihan terdistribusi dengan TensorFlow - Cara menggunakan strategi distribusi dan tautan ke banyak contoh yang menunjukkan praktik terbaik.
- Model Resmi TensorFlow - Contoh model TensorFlow 2.x canggih yang kompatibel dengan Cloud TPU.
- Panduan Performa Google Cloud TPU - Tingkatkan performa Cloud TPU lebih lanjut dengan menyesuaikan parameter konfigurasi Cloud TPU untuk aplikasi Anda.