 Xem trên TensorFlow.org |  Chạy trong Google Colab |  Xem nguồn trên GitHub |  Tải xuống sổ tay |
Hỗ trợ thử nghiệm cho TPU đám mây hiện có sẵn cho Keras và Google Colab. Trước khi bạn chạy máy tính xách tay Colab này, hãy đảm bảo rằng trình tăng tốc phần cứng của bạn là TPU bằng cách kiểm tra cài đặt máy tính xách tay của bạn: Thời gian chạy> Thay đổi loại thời gian chạy> Trình tăng tốc phần cứng> TPU.
Thiết lập
import tensorflow as tf
import os
import tensorflow_datasets as tfds
Khởi tạo TPU
TPU thường nằm trên Cloud TPU worker khác với quy trình cục bộ chạy chương trình python người dùng. Do đó, một số công việc khởi tạo cần được thực hiện để kết nối với cụm từ xa và khởi tạo TPU. Lưu ý rằng tpu lập luận để TPUClusterResolver là một địa chỉ đặc biệt chỉ dành riêng cho Colab. Trong trường hợp bạn đang chạy trên Google Compute Engine (GCE), thay vào đó bạn nên chuyển tên CloudTPU của mình.
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.58:8470 INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.58: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')]
Vị trí thiết bị thủ công
Sau khi TPU được khởi tạo, bạn có thể sử dụng vị trí thiết bị thủ công để đặt tính toán trên một thiết bị TPU duy nhất.
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)
Các chiến lược phân phối
Hầu hết người dùng muốn chạy mô hình trên nhiều TPU theo cách song song dữ liệu. Chiến lược phân phối là một sự trừu tượng có thể được sử dụng để thúc đẩy các mô hình trên CPU, GPU hoặc TPU. Chỉ cần hoán đổi chiến lược phân phối và mô hình sẽ chạy trên thiết bị nhất định. Xem hướng dẫn chiến lược phân phối để biết thêm thông tin.
Đầu tiên, tạo đối tượng 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)
Để sao chép một tính toán để nó có thể chạy trong tất cả các lõi TPU, bạn có thể chỉ cần chuyển nó đến API strategy.run . Dưới đây là một ví dụ mà tất cả các lõi sẽ nhận được các đầu vào giống nhau (a, b) và thực hiện matmul trên mỗi lõi một cách độc lập. Kết quả đầu ra sẽ là các giá trị từ tất cả các bản sao.
@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)
}
Phân loại TPU
Khi chúng ta đã học các khái niệm cơ bản, đã đến lúc xem xét một ví dụ cụ thể hơn. Hướng dẫn này trình bày cách sử dụng chiến lược phân phối tf.distribute.TPUStrategy để thúc đẩy TPU đám mây và đào tạo mô hình Keras.
Xác định mô hình Keras
Dưới đây là định nghĩa của mô hình MNIST sử dụng Keras, không thay đổi so với những gì bạn sẽ sử dụng trên CPU hoặc GPU. Lưu ý rằng việc tạo mô hình Keras cần phải ở bên trong strategy.scope , vì vậy các biến có thể được tạo trên mỗi thiết bị TPU. Các phần khác của mã không cần thiết phải nằm trong phạm vi chiến lược.
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)])
Bộ dữ liệu đầu vào
Sử dụng hiệu quả APItf.data.Dataset là rất quan trọng khi sử dụng Cloud TPU, vì không thể sử dụng Cloud TPU trừ khi bạn có thể cung cấp dữ liệu đủ nhanh cho chúng. Xem Hướng dẫn hiệu suất đường ống đầu vào để biết chi tiết về hiệu suất tập dữ liệu.
Đối với tất cả, trừ thử nghiệm đơn giản nhất (sử dụng tf.data.Dataset.from_tensor_slices hoặc dữ liệu trong biểu đồ khác), bạn sẽ cần lưu trữ tất cả các tệp dữ liệu được Dataset đọc trong nhóm Google Cloud Storage (GCS).
Đối với hầu hết các trường hợp sử dụng, bạn nên chuyển đổi dữ liệu của mình sang định dạng TFRecord và sử dụng tf.data.TFRecordDataset để đọc. Xem hướng dẫn TFRecord và tf.Example để biết chi tiết về cách thực hiện việc này. Tuy nhiên, đây không phải là một yêu cầu khó và bạn có thể sử dụng các trình đọc tập dữ liệu khác ( FixedLengthRecordDataset hoặc TextLineDataset ) nếu muốn.
Các tập dữ liệu nhỏ có thể được tải hoàn toàn vào bộ nhớ bằng tf.data.Dataset.cache .
Bất kể định dạng dữ liệu được sử dụng là gì, chúng tôi khuyên bạn nên sử dụng các tệp lớn, có dung lượng 100MB. Điều này đặc biệt quan trọng trong cài đặt nối mạng này vì chi phí mở tệp cao hơn đáng kể.
Ở đây, bạn nên sử dụng mô-đun tensorflow_datasets để nhận bản sao của dữ liệu đào tạo MNIST. Lưu ý rằng try_gcs được chỉ định để sử dụng bản sao có sẵn trong nhóm GCS công khai. Nếu bạn không chỉ định điều này, TPU sẽ không thể truy cập vào dữ liệu được tải xuống.
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
Đào tạo một mô hình bằng cách sử dụng các API cấp cao của Keras
Bạn có thể đào tạo một mô hình đơn giản với các API phù hợp / biên dịch Keras. Không có gì ở đây là TPU cụ thể, bạn sẽ viết mã tương tự bên dưới nếu bạn có nhiều GPU khác nhau và sử dụng MirroredStrategy thay vì TPUStrategy . Để tìm hiểu thêm, hãy xem hướng dẫn Phân tán với 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.3140 - sparse_categorical_accuracy: 0.8968 - val_loss: 0.0428 - val_sparse_categorical_accuracy: 0.9852 Epoch 2/5 300/300 [==============================] - 8s 26ms/step - loss: 0.0399 - sparse_categorical_accuracy: 0.9876 - val_loss: 0.0414 - val_sparse_categorical_accuracy: 0.9865 Epoch 3/5 300/300 [==============================] - 7s 25ms/step - loss: 0.0204 - sparse_categorical_accuracy: 0.9934 - val_loss: 0.0488 - val_sparse_categorical_accuracy: 0.9864 Epoch 4/5 300/300 [==============================] - 8s 25ms/step - loss: 0.0141 - sparse_categorical_accuracy: 0.9956 - val_loss: 0.0674 - val_sparse_categorical_accuracy: 0.9819 Epoch 5/5 300/300 [==============================] - 8s 25ms/step - loss: 0.0104 - sparse_categorical_accuracy: 0.9965 - val_loss: 0.0449 - val_sparse_categorical_accuracy: 0.9885 <tensorflow.python.keras.callbacks.History at 0x7fede8002940>
Để giảm python trên cao, và tối đa hóa hiệu suất của TPU của bạn, hãy thử nghiệm experimental_steps_per_execution lập luận để Model.compile . Ở đây nó tăng thông lượng lên khoảng 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 48ms/step - loss: 0.2434 - sparse_categorical_accuracy: 0.9269 - val_loss: 0.0461 - val_sparse_categorical_accuracy: 0.9861 Epoch 2/5 300/300 [==============================] - 5s 16ms/step - loss: 0.0378 - sparse_categorical_accuracy: 0.9884 - val_loss: 0.0394 - val_sparse_categorical_accuracy: 0.9877 Epoch 3/5 300/300 [==============================] - 5s 16ms/step - loss: 0.0207 - sparse_categorical_accuracy: 0.9935 - val_loss: 0.0484 - val_sparse_categorical_accuracy: 0.9861 Epoch 4/5 300/300 [==============================] - 5s 16ms/step - loss: 0.0145 - sparse_categorical_accuracy: 0.9950 - val_loss: 0.0406 - val_sparse_categorical_accuracy: 0.9878 Epoch 5/5 300/300 [==============================] - 5s 16ms/step - loss: 0.0105 - sparse_categorical_accuracy: 0.9967 - val_loss: 0.0528 - val_sparse_categorical_accuracy: 0.9857 <tensorflow.python.keras.callbacks.History at 0x7fe9242485f8>
Đào tạo một mô hình bằng cách sử dụng vòng lặp đào tạo tùy chỉnh.
Bạn cũng có thể tạo và đào tạo mô hình của mình bằng cách sử dụng trực tiếp các API tf.distribute . tf.function và tf.distribute . strategy.experimental_distribute_datasets_from_function lược.experimental_distribute_datasets_from_ Chức năng API được sử dụng để phân phối tập dữ liệu cho một hàm tập dữ liệu. Lưu ý rằng kích thước lô được chuyển vào tập dữ liệu sẽ là kích thước lô bản sao thay vì kích thước lô toàn cầu trong trường hợp này. Để tìm hiểu thêm, hãy xem hướng dẫn đào tạo Tùy chỉnh với tf.distribute.Strategy .
Đầu tiên, tạo mô hình, tập dữ liệu và các hàm tf.f.
# 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
Sau đó chạy vòng lặp đào tạo.
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.1293, accuracy: 95.98% Epoch: 1/5 Current step: 600, training loss: 0.034, accuracy: 98.94% Epoch: 2/5 Current step: 900, training loss: 0.0182, accuracy: 99.39% Epoch: 3/5 Current step: 1200, training loss: 0.0143, accuracy: 99.53% Epoch: 4/5 Current step: 1500, training loss: 0.009, accuracy: 99.69%
Cải thiện hiệu suất bằng nhiều bước trong tf.function
Hiệu suất có thể được cải thiện bằng cách chạy nhiều bước trong một tf.function . Điều này đạt được bằng cách gói lệnh gọi strategy.run với một tf.range bên trong tf.function , AutoGraph sẽ chuyển nó thành tf.while_loop trên TPU worker.
Mặc dù với hiệu suất tốt hơn, vẫn có những đánh đổi so với một bước duy nhất bên trong tf.function . Chạy nhiều bước trong một tf.function ít linh hoạt hơn, bạn không thể chạy mọi thứ một cách háo hức hoặc tùy ý mã python trong các bước.
@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.0066, accuracy: 99.78%
Bước tiếp theo
- Tài liệu Google Cloud TPU - Thiết lập và chạy Google Cloud TPU.
- Đào tạo phân tán với TensorFlow - Cách sử dụng chiến lược phân phối và liên kết đến nhiều ví dụ hiển thị các phương pháp hay nhất.
- Lưu / Tải mô hình với TensorFlow - Cách lưu và tải mô hình với các chiến lược phân phối.
- Mô hình chính thức của TensorFlow - Ví dụ về mô hình TensorFlow 2.x hiện đại tương thích với Cloud TPU.
- Hướng dẫn về hiệu suất của Google Cloud TPU - Nâng cao hơn nữa hiệu suất của Cloud TPU bằng cách điều chỉnh các thông số cấu hình Cloud TPU cho ứng dụng của bạn.