 Lihat di TensorFlow.org |  Jalankan di Google Colab |  Lihat sumber di GitHub |  Unduh buku catatan |
Sebelum Anda menjalankan notebook CoLab ini, pastikan bahwa akselerator hardware Anda adalah TPU dengan memeriksa pengaturan notebook Anda: Runtime> Ubah jenis runtime> Hardware akselerator> 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.7) or chardet (2.3.0)/charset_normalizer (2.0.8) 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 tpu argumen untuk 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.2:8470 INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.2: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 menunjukkan ini, membuat tf.distribute.TPUStrategy objek:
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 meniru perhitungan sehingga dapat berjalan di semua core TPU, Anda dapat melewati ke dalam strategy.run API. Berikut adalah contoh yang menunjukkan semua core 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 bagaimana menggunakan distribusi Strategi- tf.distribute.TPUStrategy -untuk melatih model Keras pada Cloud TPU.
Tentukan model Keras
Mulailah dengan definisi dari Sequential Model Keras untuk klasifikasi gambar pada dataset MNIST menggunakan Keras. Tidak ada bedanya dengan apa yang akan Anda gunakan jika Anda berlatih CPU atau GPU. Perhatikan bahwa kebutuhan Model Keras penciptaan berada di dalam strategy.scope , sehingga variabel dapat dibuat pada 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 yang efisien dari tf.data.Dataset API sangat penting ketika menggunakan Cloud TPU, karena tidak mungkin untuk menggunakan Cloud TPUs kecuali jika Anda dapat memberi makan data yang mereka cukup cepat. Anda dapat mempelajari lebih lanjut tentang dataset kinerja dalam panduan kinerja pipa Masukan .
Untuk semua tapi percobaan sederhana (menggunakan tf.data.Dataset.from_tensor_slices atau data lain dalam grafik), Anda perlu untuk menyimpan semua file data yang dibaca oleh Dataset di Google Cloud Storage (GCS) ember.
Untuk sebagian besar kasus penggunaan, dianjurkan untuk mengkonversi data ke dalam TFRecord dan format menggunakan tf.data.TFRecordDataset untuk membacanya. Periksa TFRecord dan tf.Example tutorial untuk rincian tentang bagaimana untuk melakukan hal ini. Ini bukan persyaratan keras dan Anda dapat menggunakan pembaca dataset lain, seperti tf.data.FixedLengthRecordDataset atau tf.data.TextLineDataset .
Anda dapat memuat seluruh dataset 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 ditunjukkan dalam kode di bawah ini, Anda harus menggunakan tensorflow_datasets modul untuk mendapatkan salinan pelatihan MNIST dan data uji. Perhatikan bahwa try_gcs ditentukan dengan menggunakan salinan yang tersedia dalam GCS ember 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 TPU spesifik dalam langkah-Anda menulis kode seperti jika Anda menggunakan GPU mutliple dan MirroredStrategy bukan TPUStrategy . Anda dapat mempelajari lebih lanjut di pelatihan Terdistribusi dengan Keras tutorial.
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.1291 - sparse_categorical_accuracy: 0.9590 - val_loss: 0.0410 - val_sparse_categorical_accuracy: 0.9871 Epoch 2/5 300/300 [==============================] - 6s 22ms/step - loss: 0.0335 - sparse_categorical_accuracy: 0.9892 - val_loss: 0.0492 - val_sparse_categorical_accuracy: 0.9847 Epoch 3/5 300/300 [==============================] - 6s 21ms/step - loss: 0.0205 - sparse_categorical_accuracy: 0.9941 - val_loss: 0.0457 - val_sparse_categorical_accuracy: 0.9868 Epoch 4/5 300/300 [==============================] - 7s 22ms/step - loss: 0.0115 - sparse_categorical_accuracy: 0.9960 - val_loss: 0.0428 - val_sparse_categorical_accuracy: 0.9891 Epoch 5/5 300/300 [==============================] - 6s 21ms/step - loss: 0.0110 - sparse_categorical_accuracy: 0.9965 - val_loss: 0.0592 - val_sparse_categorical_accuracy: 0.9853 <keras.callbacks.History at 0x7f25fc108a58>
Untuk mengurangi Python overhead dan memaksimalkan kinerja TPU Anda, lulus di berdebat- steps_per_execution -untuk 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 [==============================] - 13s 42ms/step - loss: 0.1343 - sparse_categorical_accuracy: 0.9594 - val_loss: 0.0487 - val_sparse_categorical_accuracy: 0.9837 Epoch 2/5 300/300 [==============================] - 3s 10ms/step - loss: 0.0348 - sparse_categorical_accuracy: 0.9891 - val_loss: 0.0401 - val_sparse_categorical_accuracy: 0.9876 Epoch 3/5 300/300 [==============================] - 3s 10ms/step - loss: 0.0187 - sparse_categorical_accuracy: 0.9940 - val_loss: 0.0517 - val_sparse_categorical_accuracy: 0.9844 Epoch 4/5 300/300 [==============================] - 3s 10ms/step - loss: 0.0123 - sparse_categorical_accuracy: 0.9963 - val_loss: 0.0506 - val_sparse_categorical_accuracy: 0.9860 Epoch 5/5 300/300 [==============================] - 3s 10ms/step - loss: 0.0080 - sparse_categorical_accuracy: 0.9974 - val_loss: 0.0456 - val_sparse_categorical_accuracy: 0.9878 <keras.callbacks.History at 0x7f23ac7782b0>
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 strategy.experimental_distribute_datasets_from_function API untuk mendistribusikan dataset diberikan fungsi dataset. 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, memeriksa pelatihan Kustom dengan tf.distribute.Strategy tutorial.
Pertama, buat model, set 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.1436, accuracy: 95.42% Epoch: 1/5 Current step: 600, training loss: 0.0342, accuracy: 98.94% Epoch: 2/5 Current step: 900, training loss: 0.0205, accuracy: 99.36% Epoch: 3/5 Current step: 1200, training loss: 0.0123, accuracy: 99.59% Epoch: 4/5 Current step: 1500, training loss: 0.0089, accuracy: 99.68%
Meningkatkan kinerja dengan beberapa langkah dalam tf.function
Anda dapat meningkatkan kinerja dengan menjalankan beberapa langkah dalam tf.function . Hal ini dicapai dengan membungkus strategy.run panggilan dengan tf.range dalam tf.function , dan tanda tangan akan mengubahnya menjadi tf.while_loop pada pekerja TPU.
Meskipun peningkatan kinerja, ada timbal balik dengan metode ini dibandingkan dengan berjalan satu langkah dalam tf.function . Menjalankan beberapa langkah dalam tf.function kurang fleksibel-Anda tidak dapat menjalankan hal-hal bersemangat atau kode Python sewenang-wenang dalam langkah-langkah.
@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.0093, accuracy: 99.68%
Langkah selanjutnya
- Dokumentasi Google Cloud TPU : Cara mengatur dan menjalankan Google Cloud TPU.
- Notebook Google Cloud TPU CoLab : End-to-end contoh pelatihan.
- Panduan kinerja Google Cloud TPU : Meningkatkan kinerja Cloud TPU lanjut dengan menyesuaikan parameter konfigurasi Cloud TPU untuk aplikasi Anda
- Didistribusikan pelatihan dengan TensorFlow : Cara menggunakan strategi-termasuk distribusi
tf.distribute.TPUStrategycontoh -dengan menunjukkan praktik terbaik.