Google I / O 18-20 मई को लौटता है! जगह आरक्षित करें और अपना शेड्यूल बनाएं अभी रजिस्टर करें
इस पेज का अनुवाद Cloud Translation API से किया गया है.
Switch to English

TPU का प्रयोग करें

TensorFlow.org पर देखें Google Colab में चलाएं GitHub पर स्रोत देखें नोटबुक डाउनलोड करें

क्लाउड टीपीयू के लिए प्रायोगिक समर्थन वर्तमान में केरस और Google कोलाब के लिए उपलब्ध है। इससे पहले कि आप इस Colab नोटबुक को चलाएं, सुनिश्चित करें कि आपकी हार्डवेयर त्वरक आपकी नोटबुक सेटिंग्स की जाँच करके एक TPU है: रनटाइम> रनटाइम प्रकार बदलें> हार्डवेयर त्वरक> TPU।

सेट अप

import tensorflow as tf

import os
import tensorflow_datasets as tfds

टीपीयू इनिशियलाइज़ेशन

टीपीयू आमतौर पर क्लाउड टीपीयू श्रमिकों पर होते हैं जो उपयोगकर्ता अजगर कार्यक्रम चलाने वाली स्थानीय प्रक्रिया से भिन्न होते हैं। इस प्रकार दूरस्थ क्लस्टर से जुड़ने और TPU को आरंभ करने के लिए कुछ आरंभिक कार्य किए जाने की आवश्यकता है। ध्यान दें कि TPUClusterResolver का tpu तर्क केवल TPUClusterResolver लिए एक विशेष पता है। उस स्थिति में जब आप Google कंप्यूट इंजन (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')]

मैनुअल डिवाइस प्लेसमेंट

टीपीयू के आरंभ होने के बाद, आप एक टीपीयू डिवाइस पर कम्प्यूटेशन लगाने के लिए मैनुअल डिवाइस प्लेसमेंट का उपयोग कर सकते हैं।

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)

वितरण रणनीतियों

अधिकांश बार उपयोगकर्ता डेटा को समानांतर रूप से कई टीपीयू पर मॉडल चलाना चाहते हैं। एक वितरण रणनीति एक अमूर्त है जो सीपीयू, जीपीयू या टीपीयू पर मॉडल को चलाने के लिए इस्तेमाल किया जा सकता है। बस वितरण रणनीति को स्वैप करें और मॉडल दिए गए डिवाइस पर चलेगा। अधिक जानकारी के लिए वितरण रणनीति गाइड देखें।

सबसे पहले, 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)

एक संगणना को दोहराने के लिए यह सभी टीपीयू कोर में चल सकता है, आप बस इसे 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 मॉडल को प्रशिक्षित करने के लिए वितरण रणनीति tf.distribute.TPUStrategy का उपयोग कैसे करें।

एक केरस मॉडल को परिभाषित करें

नीचे Keras का उपयोग करके MNIST मॉडल की परिभाषा दी गई है, जो आप CPU या GPU पर उपयोग करेंगे से अपरिवर्तित है। ध्यान दें कि केरस मॉडल निर्माण के लिए 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 का उपयोग करते समयtf.data.Dataset API का कुशल उपयोग महत्वपूर्ण है, क्योंकि क्लाउडtf.data.Dataset का उपयोग करना असंभव है जब तक कि आप उन्हें डेटा जल्दी से पर्याप्त फीड नहीं कर सकते। डेटासेट परफॉर्मेंस के विवरण के लिए इनपुट पाइपलाइन प्रदर्शन गाइड देखें।

सभी के लिए लेकिन सबसे सरल प्रयोग ( tf.data.Dataset.from_tensor_slices या अन्य-इन-ग्राफ डेटा का उपयोग करके) आपको डेटा क्लाउड द्वारा पढ़ी गई सभी डेटा फ़ाइलों को Google क्लाउड स्टोरेज (GCS): बाल्टी में संग्रहीत करना होगा।

अधिकांश उपयोग-मामलों के लिए, अपने डेटा को TFRecord प्रारूप में परिवर्तित करने और इसे पढ़ने के लिए tf.data.TFRecordDataset का उपयोग करने की tf.data.TFRecordDataset जाती है। यह कैसे करें के विवरण के लिए TFRecord और tf.Example ट्यूटोरियल देखें। यह, हालांकि, एक कठिन आवश्यकता नहीं है और यदि आप चाहें तो आप अन्य डेटासेट रीडर ( FixedLengthRecordDataset या TextLineDataset ) का उपयोग कर सकते हैं।

छोटे डेटासेट को पूरी तरह से tf.data.Dataset.cache का उपयोग करके मेमोरी में लोड किया जा सकता है।

डेटा प्रारूप का उपयोग किए बिना, यह दृढ़ता से अनुशंसा की जाती है कि आप 100 एमबी के आदेश पर बड़ी फ़ाइलों का उपयोग करें। यह इस नेटवर्क सेटिंग में विशेष रूप से महत्वपूर्ण है क्योंकि फ़ाइल खोलने का ओवरहेड काफी अधिक है।

यहां आपको MNIST प्रशिक्षण डेटा की एक प्रति प्राप्त करने के लिए tensorflow_datasets मॉड्यूल का उपयोग करना चाहिए। ध्यान दें कि try_gcs एक सार्वजनिक जीसीएस बाल्टी में उपलब्ध प्रतिलिपि का उपयोग करने के लिए निर्दिष्ट है। यदि आप इसे निर्दिष्ट नहीं करते हैं, तो टीपीयू डाउनलोड किए गए डेटा तक नहीं पहुंच पाएगा।

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

केरस उच्च स्तरीय एपीआई का उपयोग करके एक मॉडल को प्रशिक्षित करें

आप बस केरेस फिट / संकलन एपीआई के साथ एक मॉडल को प्रशिक्षित कर सकते हैं। यहां कुछ नहीं है, TPU विशिष्ट है कि आप नीचे अगर तुम mutliple GPUs था और जहां एक का उपयोग कर एक ही कोड लिखते थे MirroredStrategy बजाय एक TPUStrategy । अधिक जानने के लिए, केरेस ट्यूटोरियल के साथ वितरित प्रशिक्षण देखें

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)
075bdb690

अजगर ओवरहेड को कम करने के लिए, और अपने TPU के प्रदर्शन को अधिकतम करने के लिए, Model.compile लिए प्रयोगात्मक 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 APIs का उपयोग करके भी बना और प्रशिक्षित कर सकते हैं। strategy.experimental_distribute_datasets_from_function एपीआई को एक डाटासेट फंक्शन दिए गए डेटा को वितरित करने के लिए उपयोग किया जाता है। ध्यान दें कि डेटासेट में पारित बैच आकार इस मामले में वैश्विक बैच आकार के बजाय प्रति बैच आकार के अनुसार होगा। अधिक जानने के लिए, 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 भीतर कई चरणों को चलाकर प्रदर्शन में सुधार किया जा सकता है। यह tf.range अंदर tf.function साथ strategy.run कॉल को लपेटकर प्राप्त किया जाता है, AutoGraph इसे TPU कार्यकर्ता पर tf.while_loop बदल देगा।

हालांकि बेहतर प्रदर्शन के साथ, tf.function अंदर एक एकल चरण के साथ तुलना करने वाले tf.functiontf.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%

अगला कदम