استخدم TPU

عرض على TensorFlow.org تشغيل في Google Colab عرض المصدر على جيثب تحميل دفتر

قبل تشغيل الكمبيوتر الدفتري Colab هذا ، تأكد من أن مسرع أجهزتك عبارة عن TPU من خلال التحقق من إعدادات الكمبيوتر الدفتري: وقت التشغيل > تغيير نوع وقت التشغيل > مسرع الأجهزة > TPU .

يثبت

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)

تهيئة TPU

عادةً ما تكون شركات TPU عمال Cloud TPU ، والتي تختلف عن العملية المحلية التي تقوم بتشغيل برنامج Python الخاص بالمستخدم. وبالتالي ، تحتاج إلى القيام ببعض أعمال التهيئة للاتصال بالمجموعة البعيدة وتهيئة TPU. لاحظ أن الوسيطة tpu لـ tf.distribute.cluster_resolver.TPUClusterResolver هي عنوان خاص لـ Colab فقط. إذا كنت تقوم بتشغيل الكود الخاص بك على Google Compute Engine (GCE) ، فيجب عليك بدلاً من ذلك تمرير اسم Cloud TPU الخاص بك.

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

وضع الجهاز اليدوي

بعد تهيئة 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 بطريقة موازية للبيانات. لتوزيع نموذجك على العديد من TPU (أو مسرعات أخرى) ، يقدم TensorFlow العديد من استراتيجيات التوزيع. يمكنك استبدال إستراتيجية التوزيع الخاصة بك وسيعمل النموذج على أي جهاز (TPU) محدد. تحقق من دليل استراتيجية التوزيع لمزيد من المعلومات.

لتوضيح ذلك ، قم بإنشاء كائن 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)

لتكرار عملية حسابية بحيث يمكن تشغيلها في جميع نوى TPU ، يمكنك تمريرها إلى strategy.run API. يوجد أدناه مثال يوضح أن جميع النوى تتلقى نفس المدخلات (a, b) وتقوم بضرب المصفوفة على كل نواة بشكل مستقل. ستكون المخرجات هي القيم من جميع النسخ المتماثلة.

@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 - لتدريب نموذج Keras على Cloud TPU.

تحديد نموذج Keras

ابدأ بتعريف نموذج Keras Sequential لتصنيف الصور في مجموعة بيانات MNIST باستخدام Keras. لا يختلف الأمر عما ستستخدمه إذا كنت تتدرب على وحدات المعالجة المركزية أو وحدات معالجة الرسومات. لاحظ أن إنشاء نموذج 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)])

قم بتحميل مجموعة البيانات

يعد الاستخدام الفعال لـ tf.data.Dataset API أمرًا بالغ الأهمية عند استخدام Cloud TPU ، لأنه من المستحيل استخدام Cloud TPUs إلا إذا كان بإمكانك تزويدهم بالبيانات بسرعة كافية. يمكنك معرفة المزيد حول أداء مجموعة البيانات في دليل أداء مسار الإدخال .

بالنسبة لجميع التجارب باستثناء أبسطها (باستخدام tf.data.Dataset.from_tensor_slices أو بيانات أخرى في الرسم البياني) ، تحتاج إلى تخزين جميع ملفات البيانات التي قرأتها مجموعة البيانات في حاويات Google Cloud Storage (GCS).

بالنسبة لمعظم حالات الاستخدام ، يوصى بتحويل بياناتك إلى تنسيق TFRecord واستخدام tf.data.TFRecordDataset لقراءتها. تحقق من TFRecord و tf مثال تعليمي للحصول على تفاصيل حول كيفية القيام بذلك. إنه ليس مطلبًا صعبًا ويمكنك استخدام قارئات مجموعة بيانات أخرى ، مثل tf.data.FixedLengthRecordDataset أو tf.data.TextLineDataset .

يمكنك تحميل مجموعات بيانات صغيرة كاملة في الذاكرة باستخدام 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)

  # 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

تدريب النموذج باستخدام واجهات برمجة تطبيقات Keras عالية المستوى

يمكنك تدريب النموذج الخاص بك باستخدام Keras fit and compile APIs. لا يوجد شيء خاص بـ TPU في هذه الخطوة - تكتب الكود كما لو كنت تستخدم وحدات معالجة رسومات متعددة واستراتيجية 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 [==============================] - 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>

لتقليل حمل Python وتعظيم أداء TPU الخاص بك ، قم بتمرير الوسيطة - 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.
                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>

تدريب النموذج باستخدام حلقة تدريب مخصصة

يمكنك أيضًا إنشاء نموذجك وتدريبه باستخدام tf.function و tf.distribute APIs مباشرة. يمكنك استخدام strategy.experimental_distribute_datasets_from_function API لتوزيع مجموعة البيانات المعطاة لوظيفة مجموعة البيانات. لاحظ أنه في المثال أدناه ، يكون حجم الدُفعة الذي تم تمريره إلى مجموعة البيانات هو حجم الدُفعة لكل نسخة متماثلة بدلاً من حجم الدُفعة العام. لمعرفة المزيد ، تحقق من التدريب المخصص باستخدام tf.distribute.Strategy .

أولاً ، قم بإنشاء النموذج ومجموعات البيانات ووظائف tf:

# 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

ثم قم بتشغيل حلقة التدريب:

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%

تحسين الأداء بخطوات متعددة داخل tf.function

يمكنك تحسين الأداء عن طريق تنفيذ خطوات متعددة داخل tf.function . يتم تحقيق ذلك عن طريق التفاف استدعاء tf.range strategy.run tf.function ، وسيقوم AutoGraph بتحويلها إلى tf. tf.while_loop على عامل TPU.

على الرغم من الأداء المحسن ، إلا أن هناك مفاضلات مع هذه الطريقة مقارنة بتشغيل خطوة واحدة داخل tf.function . يعد تشغيل عدة خطوات في tf.function أقل مرونة - لا يمكنك تشغيل الأشياء بفارغ الصبر أو تشغيل كود Python العشوائي ضمن الخطوات.

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

الخطوات التالية