Lihat di TensorFlow.org | Jalankan di Google Colab | Lihat sumber di GitHub | Unduh buku catatan |
Sebelum menjalankan notebook Colab ini, pastikan akselerator perangkat keras Anda adalah TPU dengan memeriksa pengaturan notebook Anda: Waktu proses > Ubah jenis waktu proses > Akselerator perangkat keras > TPU .
Mempersiapkan
import tensorflow as tf
import os
import tensorflow_datasets as tfds
/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)
inisialisasi TPU
TPU biasanya adalah pekerja Cloud TPU, yang berbeda dari proses lokal yang menjalankan program Python pengguna. Dengan demikian, Anda perlu melakukan beberapa pekerjaan inisialisasi untuk terhubung ke kluster jarak jauh dan menginisialisasi TPU. Perhatikan bahwa argumen tpu
ke tf.distribute.cluster_resolver.TPUClusterResolver
adalah alamat khusus hanya untuk Colab. Jika Anda menjalankan kode di Google Compute Engine (GCE), Anda harus meneruskan nama Cloud TPU 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: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')]
Penempatan perangkat manual
Setelah TPU diinisialisasi, Anda dapat menggunakan penempatan perangkat manual untuk menempatkan perhitungan 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
Biasanya Anda menjalankan model Anda pada beberapa TPU dengan cara paralel data. Untuk mendistribusikan model Anda di beberapa TPU (atau akselerator lainnya), TensorFlow menawarkan beberapa strategi distribusi. Anda dapat mengganti strategi distribusi Anda dan model akan berjalan pada perangkat (TPU) apa pun. Periksa panduan strategi distribusi untuk informasi lebih lanjut.
Untuk mendemonstrasikan ini, buat objek tf.distribute.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 sehingga dapat berjalan di semua inti TPU, Anda dapat meneruskannya ke API strategy.run
. Di bawah ini adalah contoh yang menunjukkan semua inti menerima input yang sama (a, b)
dan melakukan perkalian matriks pada setiap inti secara independen. Outputnya akan menjadi 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
Setelah membahas konsep dasar, pertimbangkan contoh yang lebih konkret. Bagian ini menunjukkan cara menggunakan strategi distribusi— tf.distribute.TPUStrategy
—untuk melatih model Keras di Cloud TPU.
Tentukan model Keras
Mulailah dengan definisi model Keras Sequential
untuk klasifikasi gambar pada dataset MNIST menggunakan Keras. Ini tidak berbeda dengan apa yang akan Anda gunakan jika Anda berlatih 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 lingkup 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)])
Muat kumpulan data
Penggunaan tf.data.Dataset
API secara efisien sangat penting saat menggunakan Cloud TPU, karena Cloud TPU tidak mungkin digunakan kecuali Anda dapat memberi mereka data dengan cukup cepat. Anda dapat mempelajari lebih lanjut tentang kinerja kumpulan data dalam panduan kinerja saluran masukan .
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 Dataset 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 tutorial TFRecord dan tf.Example untuk detail tentang cara melakukannya. Ini bukan persyaratan yang sulit dan Anda dapat menggunakan pembaca kumpulan data lain, seperti tf.data.FixedLengthRecordDataset
atau tf.data.TextLineDataset
.
Anda dapat memuat seluruh kumpulan data kecil ke dalam memori menggunakan tf.data.Dataset.cache
.
Terlepas dari format data yang digunakan, sangat disarankan agar Anda menggunakan file besar di urutan 100MB. Ini sangat penting dalam pengaturan jaringan ini, karena biaya pembukaan file secara signifikan lebih tinggi.
Seperti yang ditunjukkan pada kode di bawah ini, Anda harus menggunakan modul tensorflow_datasets
untuk mendapatkan salinan data pelatihan dan pengujian MNIST. Perhatikan bahwa try_gcs
ditentukan untuk menggunakan salinan yang tersedia di bucket GCS publik. Jika Anda tidak menentukan ini, 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)
# 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
Latih model menggunakan API tingkat tinggi Keras
Anda dapat melatih model Anda dengan Keras fit
dan compile
API. Tidak ada yang spesifik untuk TPU dalam langkah ini—Anda menulis kode seolah-olah Anda menggunakan beberapa GPU dan MirroredStrategy
alih-alih TPUStrategy
. Anda dapat mempelajari lebih lanjut di pelatihan Terdistribusi dengan tutorial 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 [==============================] - 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>
Untuk mengurangi overhead Python dan memaksimalkan kinerja TPU Anda, berikan argumen— steps_per_execution
—ke Model.compile
. Dalam contoh ini, ini meningkatkan throughput sekitar 50%:
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)
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>
Latih model menggunakan loop pelatihan khusus
Anda juga dapat membuat dan melatih model Anda menggunakan tf.function
dan tf.distribute
API secara langsung. Anda dapat menggunakan API strategy.experimental_distribute_datasets_from_function
untuk mendistribusikan kumpulan data yang diberikan fungsi kumpulan data. Perhatikan bahwa dalam contoh di bawah ukuran batch yang diteruskan ke dataset adalah ukuran batch per replika, bukan ukuran batch global. Untuk mempelajari lebih lanjut, lihat pelatihan Kustom dengan tutorial tf.distribute.Strategy .
Pertama, buat model, kumpulan data, dan tf.functions:
# 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),))
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
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.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%
Meningkatkan kinerja dengan beberapa langkah di dalam tf.function
Anda dapat meningkatkan kinerja dengan menjalankan beberapa langkah dalam tf.function
. Ini dicapai dengan membungkus panggilan strategy.run
dengan tf.range
di dalam tf.function
, dan AutoGraph akan mengonversinya menjadi tf.while_loop
pada pekerja TPU.
Terlepas dari peningkatan kinerja, ada pengorbanan dengan metode ini dibandingkan dengan menjalankan satu langkah di dalam tf.function
. Menjalankan beberapa langkah dalam tf.function
kurang fleksibel—Anda tidak dapat menjalankan sesuatu dengan bersemangat atau kode Python arbitrer dalam langkah-langkahnya.
@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.0081, accuracy: 99.74%
Langkah selanjutnya
- Dokumentasi Google Cloud TPU : Cara menyiapkan dan menjalankan Google Cloud TPU.
- Notebook Google Cloud TPU Colab : Contoh pelatihan menyeluruh.
- Panduan kinerja Google Cloud TPU : Tingkatkan kinerja Cloud TPU lebih lanjut dengan menyesuaikan parameter konfigurasi Cloud TPU untuk aplikasi Anda
- Pelatihan terdistribusi dengan TensorFlow : Cara menggunakan strategi distribusi—termasuk
tf.distribute.TPUStrategy
—dengan contoh yang menunjukkan praktik terbaik.