टीपीयू का प्रयोग करें

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

इस 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 को प्रारंभ करने के लिए कुछ आरंभीकरण कार्य करने की आवश्यकता है। ध्यान दें कि tf.distribute.cluster_resolver.TPUClusterResolver का tpu तर्क केवल tf.distribute.cluster_resolver.TPUClusterResolver के लिए एक विशेष पता है। यदि आप अपना कोड Google कंप्यूट इंजन (GCE) पर चला रहे हैं, तो आपको इसके बजाय अपने क्लाउड 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')]

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

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

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 (या अन्य त्वरक) पर वितरित करने के लिए, TensorFlow कई वितरण रणनीतियाँ प्रदान करता है। आप अपनी वितरण रणनीति को बदल सकते हैं और मॉडल किसी दिए गए (टीपीयू) डिवाइस पर चलेगा। अधिक जानकारी के लिए वितरण रणनीति मार्गदर्शिका देखें।

इसे प्रदर्शित करने के लिए, एक 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)

एक गणना को दोहराने के लिए ताकि यह सभी टीपीयू कोर में चल सके, आप इसे strategy.run एपीआई में पास कर सकते हैं। नीचे एक उदाहरण दिया गया है जो सभी कोर को समान इनपुट (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)
}

टीपीयू पर वर्गीकरण

बुनियादी अवधारणाओं को कवर करने के बाद, एक अधिक ठोस उदाहरण पर विचार करें। यह खंड दर्शाता है कि क्लाउड टीपीयू पर केरस मॉडल को प्रशिक्षित करने के लिए वितरण रणनीति- tf.distribute.TPUStrategy का उपयोग कैसे करें।

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

केरस का उपयोग करके एमएनआईएसटी डेटासेट पर छवि वर्गीकरण के लिए Sequential केरस मॉडल की परिभाषा के साथ प्रारंभ करें। यदि आप सीपीयू या जीपीयू पर प्रशिक्षण ले रहे थे तो यह आपके द्वारा उपयोग किए जाने वाले कार्यों से अलग नहीं है। ध्यान दें कि केरस मॉडल निर्माण strategy.scope के अंदर होना चाहिए, इसलिए प्रत्येक टीपीयू डिवाइस पर चर बनाए जा सकते हैं। कोड के अन्य भागों का रणनीति के दायरे में होना आवश्यक नहीं है।

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 का कुशल उपयोग महत्वपूर्ण है, क्योंकि क्लाउड टीपीयू का उपयोग करना असंभव है जब तक कि आप उन्हें पर्याप्त रूप से पर्याप्त डेटा नहीं खिला सकते। आप इनपुट पाइपलाइन प्रदर्शन मार्गदर्शिका में डेटासेट प्रदर्शन के बारे में अधिक जान सकते हैं।

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

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

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

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

जैसा कि नीचे दिए गए कोड में दिखाया गया है, आपको MNIST प्रशिक्षण और परीक्षण डेटा की एक प्रति प्राप्त करने के लिए tensorflow_datasets मॉड्यूल का उपयोग करना चाहिए। ध्यान दें कि try_gcs को सार्वजनिक 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)

  # 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

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

आप अपने मॉडल को केरस fit और compile एपीआई के साथ प्रशिक्षित कर सकते हैं। इस चरण में टीपीयू-विशिष्ट कुछ भी नहीं है—आप कोड को ऐसे लिखते हैं जैसे कि आप TPUStrategy MirroredStrategy उपयोग कर रहे थे। आप केरस ट्यूटोरियल के साथ वितरित प्रशिक्षण में अधिक जान सकते हैं।

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>

पायथन ओवरहेड को कम करने और अपने टीपीयू के प्रदर्शन को अधिकतम करने के लिए, तर्क में पास steps_per_execution —to 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 API का उपयोग करके अपना मॉडल बना और प्रशिक्षित भी कर सकते हैं। डेटासेट फ़ंक्शन दिए गए डेटासेट को वितरित करने के लिए आप strategy.experimental_distribute_datasets_from_function API का उपयोग कर सकते हैं। ध्यान दें कि डेटासेट में पारित बैच आकार के नीचे के उदाहरण में वैश्विक बैच आकार के बजाय प्रति-प्रतिकृति बैच आकार है। अधिक जानने के लिए, tf.distribute.Strategy ट्यूटोरियल के साथ कस्टम प्रशिक्षण देखें

सबसे पहले, मॉडल, डेटासेट और 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

फिर, प्रशिक्षण लूप चलाएँ:

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 के अंदर tf.function के साथ strategy.run कॉल को लपेटकर हासिल किया जाता है, और AutoGraph इसे TPU कार्यकर्ता पर tf. tf.while_loop में बदल देगा।

बेहतर प्रदर्शन के बावजूद, 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.0081, accuracy: 99.74%
प्लेसहोल्डर22

अगले कदम