![]() | ![]() | ![]() | ![]() |
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
- 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.
- Menyimpan / Memuat model dengan TensorFlow - Cara menyimpan dan memuat model dengan strategi distribusi.
- 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.