Sử dụng TPU

Xem trên TensorFlow.org Chạy trong Google Colab Xem nguồn trên GitHub Tải xuống sổ ghi chép

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 .

Thành lập

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)

Khởi tạo TPU

TPU thường là các nhân viên Cloud TPU, khác với quy trình cục bộ chạy chương trình Python của người dùng. Do đó, bạn cần thực hiện một số công việc khởi tạo để kết nối với cụm từ xa và khởi tạo TPU. Lưu ý rằng đối số tpu cho tf.distribute.cluster_resolver.TPUClusterResolver là một địa chỉ đặc biệt chỉ dành cho Colab. Nếu bạn đang chạy mã của mình trên Google Compute Engine (GCE), thay vào đó bạn nên chuyển tên của Cloud TPU của mình.

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')]

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

Thông thường, bạn chạy mô hình của mình trên nhiều TPU theo cách song song dữ liệu. Để phân phối mô hình của bạn trên nhiều TPU (hoặc các bộ tăng tốc khác), TensorFlow đưa ra một số chiến lược phân phối. Bạn có thể thay thế chiến lược phân phối của mình và mô hình sẽ chạy trên mọi thiết bị (TPU) nhất định. Kiểm tra hướng dẫn chiến lược phân phối để biết thêm thông tin.

Để chứng minh điều này, hãy tạo một đối tượng 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)

Để sao chép một phép tính để nó có thể chạy trong tất cả các lõi TPU, bạn có thể chuyển nó vào API strategy.run . Dưới đây là một ví dụ cho thấy tất cả các lõi nhận cùng một đầu vào (a, b) và thực hiện phép nhân ma trận 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 trên TPUs

Sau khi đề cập đến các khái niệm cơ bản, hãy xem xét một ví dụ cụ thể hơn. Phần này trình bày cách sử dụng chiến lược phân tf.distribute.TPUStrategy —để đào tạo mô hình Keras trên Cloud TPU.

Xác định mô hình Keras

Bắt đầu với định nghĩa về mô hình Keras Sequential để phân loại hình ảnh trên tập dữ liệu MNIST sử dụng Keras. Nó không khác gì những gì bạn sẽ sử dụng nếu bạn đang đào tạo về 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 lược.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)])

Tải tập dữ liệu

Việc sử dụng hiệu quả API tf.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. Bạn có thể tìm hiểu thêm về hiệu suất tập dữ liệu trong Hướng dẫn hiệu suất đường ống đầu vào .

Đối với tất cả các thử nghiệm, trừ những 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 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. Kiểm tra hướng dẫn TFRecord và tf.Example để biết chi tiết về cách thực hiện việc này. Đâ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, chẳng hạn như tf.data.FixedLengthRecordDataset hoặc tf.data.TextLineDataset .

Bạn có thể tải toàn bộ tập dữ liệu nhỏ 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ì, 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ể.

Như được hiển thị trong đoạn mã bên dưới, bạn nên sử dụng mô-đun tensorflow_datasets để nhận bản sao dữ liệu kiểm tra và đà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 đã 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)

  # 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

Đào tạo 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ô hình của mình bằng Keras fit và compile API. Không có gì cụ thể về TPU trong bước này — bạn viết mã như thể bạn đang sử dụng nhiều GPU khác nhau và MirroredStrategy thay vì TPUStrategy . Bạn có thể tìm hiểu thêm trong phần đào tạo Phân tán với hướng dẫn 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>

Để giảm chi phí Python và tối đa hóa hiệu suất của TPU của bạn, hãy chuyển đối steps_per_execution —to Model.compile . Trong ví dụ nà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.
                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>

Đào tạo 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 Bạn có thể sử dụng API strategy.experimental_distribute_datasets_from_function Chức năng để phân phối tập dữ liệu được cung cấp cho một hàm tập dữ liệu. Lưu ý rằng trong ví dụ dưới đây, kích thước lô được chuyển vào tập dữ liệu là kích thước lô cho mỗi bản sao thay vì kích thước lô chung. Để 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à 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

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.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%

Cải thiện hiệu suất với nhiều bước bên trong tf.function

Bạn có thể cải thiện hiệu suất 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 . function và AutoGraph sẽ chuyển nó thành tf.while_loop trên TPU worker.

Mặc dù hiệu suất được cải thiện, phương pháp này có những đánh đổi so với việc chạy 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.0081, accuracy: 99.74%

Bước tiếp theo