تاریخ را ذخیره کنید! Google I / O 18-20 مه بازمی گردد اکنون ثبت نام کنید
این صفحه به‌وسیله ‏Cloud Translation API‏ ترجمه شده است.
Switch to English

از TPU استفاده کنید

مشاهده در TensorFlow.org در Google Colab اجرا کنید مشاهده منبع در GitHub دانلود دفترچه یادداشت

پشتیبانی آزمایشی Cloud TPU در حال حاضر برای Keras و Google Colab در دسترس است. قبل از اجرای این نوت بوک های Colab ، با بررسی تنظیمات نوت بوک ، اطمینان حاصل کنید که شتاب دهنده سخت افزار شما TPU است: زمان اجرا> تغییر نوع زمان اجرا> سخت افزار سخت افزار> TPU.

برپایی

import tensorflow as tf

import os
import tensorflow_datasets as tfds

شروع TPU

TPU ها معمولاً روی کارگران Cloud TPU هستند که با روند محلی اجرای برنامه پایتون کاربر متفاوت هستند. بنابراین برای اتصال به خوشه از راه دور و مقداردهی اولیه TPU باید برخی کارهای اولیه را انجام داد. توجه داشته باشید که آرگومان tpu برای TPUClusterResolver آدرس خاصی فقط برای Colab است. درصورت استفاده از Google Compute Engine (GCE) ، باید در عوض نام CloudTPU خود را وارد کنید.

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

قرار دادن دستگاه دستی

پس از مقدار دهی اولیه TPU ، می توانید با استفاده از قرار دادن دستی دستگاه ، محاسبات را روی یک دستگاه 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)

استراتژی های توزیع

بیشتر اوقات کاربران می خواهند مدل را روی چندین TPU به صورت موازی داده اجرا کنند. استراتژی توزیع انتزاعی است که می تواند برای هدایت مدل ها روی CPU ، GPU یا TPU استفاده شود. نگران نباشید ، استراتژی توزیع را عوض کنید و مدل بر روی دستگاه داده شده اجرا شود. برای اطلاعات بیشتر به راهنمای استراتژی توزیع مراجعه کنید.

ابتدا شی 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)

به تکرار یک محاسبه بنابراین می تواند در تمام هسته TPU اجرا شود، شما می توانید به سادگی آن را به تصویب strategy.run API. در زیر مثالی آورده شده است که کلیه هسته ها ورودی های یکسانی (a, b) بدست می آورند و روی هر هسته به طور مستقل matmul را انجام می دهند. خروجی ها مقادیر حاصل از همه نسخه ها خواهند بود.

@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)
}

طبقه بندی در TPU ها

همانطور که مفاهیم اساسی را آموختیم ، وقت آن است که به یک مثال مشخص تر نگاه کنیم. این راهنما نحوه استفاده از استراتژی توزیع tf.distribute.TPUStrategy برای هدایت Cloud TPU و آموزش مدل Keras نشان می دهد.

مدل کراس را تعریف کنید

در زیر تعریف مدل MNIST با استفاده از Keras ، بدون تغییر از آنچه در پردازنده یا GPU استفاده می کنید ، آورده شده است. توجه داشته باشید که نیازهای ایجاد مدل Keras به داخل strategy.scope ، به طوری که متغیرها را می توان در هر دستگاه TPU ایجاد شده است. سایر قسمتهای کد لازم نیست که در محدوده استراتژی باشند.

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

مجموعه داده های ورودی

استفاده م oftf.data.Dataset APItf.data.Dataset هنگام استفاده از Cloud TPU بسیار حیاتی است ، زیرا استفاده از Cloud TPU غیرممکن است مگر اینکه بتوانید داده های آنها را با سرعت کافی تغذیه کنید. برای جزئیات بیشتر در مورد عملکرد مجموعه داده ها ، به راهنمای عملکرد خط لوله ورودی مراجعه کنید.

برای همه موارد ، به جز ساده ترین آزمایش (با استفاده از tf.data.Dataset.from_tensor_slices یا سایر داده های tf.data.Dataset.from_tensor_slices ) ، شما باید تمام پرونده های داده ای را که توسط Dataset در سطل ذخیره سازی Google Cloud Storage (GCS) خوانده شده ذخیره کنید.

برای بیشتر موارد استفاده ، توصیه می شود داده های خود را به قالب TFRecord تبدیل کرده و برای خواندن آنها از tf.data.TFRecordDataset استفاده کنید. برای جزئیات بیشتر در مورد چگونگی انجام این کار ، به TFRecord و tf مراجعه کنید. اگرچه این مورد سختی نیست و در صورت تمایل می توانید از خوانندگان مجموعه داده دیگر ( FixedLengthRecordDataset یا TextLineDataset ) استفاده کنید.

مجموعه داده های کوچک را می توان با استفاده از tf.data.Dataset.cache به طور کامل در حافظه tf.data.Dataset.cache .

صرف نظر از قالب داده مورد استفاده ، اکیداً توصیه می شود که از پرونده های بزرگ و به سفارش 100 مگابایت استفاده کنید. این امر به ویژه در این تنظیم شبکه بسیار مهم است زیرا سربار باز کردن یک پرونده به طور قابل توجهی بالاتر است.

در اینجا شما باید از ماژول tensorflow_datasets برای گرفتن نسخه ای از داده های آموزش MNIST استفاده کنید. توجه داشته باشید که try_gcs برای استفاده از کپی موجود در سطل عمومی GCS مشخص شده است. اگر این مورد را مشخص نکنید ، TPU نمی تواند به داده های بارگیری شده دسترسی پیدا کند.

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

با استفاده از API های سطح بالای Keras مدلی را آموزش دهید

شما می توانید به سادگی با Keras fit / compile API یک مدل را آموزش دهید. هیچ چیز در اینجا TPU خاص است، شما می توانید همان کد زیر اگر شما تا به حال GPU ها های mutliple و که در آن با استفاده از یک ارسال MirroredStrategy به جای یک TPUStrategy . برای کسب اطلاعات بیشتر ، آموزش توزیع شده با آموزش 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>

به منظور کاهش سربار پایتون، و به حداکثر رساندن عملکرد TPU خود را، سعی کنید از تجربی experimental_steps_per_execution آرگومان به Model.compile . در اینجا میزان تولید حدود 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>

با استفاده از حلقه آموزش سفارشی ، یک مدل را آموزش دهید.

همچنین می توانید با استفاده از tf.function و tf.distribute API به طور مستقیم مدل های خود را ایجاد و آموزش دهید. strategy.experimental_distribute_datasets_from_function API برای توزیع مجموعه داده با توجه به عملکرد مجموعه داده استفاده می شود. توجه داشته باشید که اندازه دسته ای که به مجموعه داده منتقل می شود ، به جای اندازه دسته جهانی در این مورد ، به ازای هر نمونه اندازه خواهد بود. برای کسب اطلاعات بیشتر ، آموزش سفارشی با tf.distribute.Strategy را ببینید .

ابتدا مدل ، مجموعه داده ها و 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

سپس حلقه آموزش را اجرا کنید.

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%

بهبود عملکرد با چندین مرحله در tf.function

با اجرای چندین مرحله در یک tf.function می توان عملکرد را بهبود بخشید. این امر با بسته بندی strategy.run با tf.range درون tf.function tf.fo به دست می آید ، AutoGraph آن را به tf.while_loop روی TPU کارگر تبدیل می کند.

اگرچه با عملکرد بهتر ، در مقایسه با یک مرحله در داخل tf.function . tf.function . اجرای چند مرحله در یک tf.function انعطاف پذیری کمتری دارد ، شما نمی توانید با اشتیاق یا خودسرانه کد پایتون را در مراحل اجرا کنید.

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

مراحل بعدی