This page was translated by the 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 ซึ่งแตกต่างจากกระบวนการภายในที่รันโปรแกรม python ของผู้ใช้ ดังนั้นจึงจำเป็นต้องดำเนินการเริ่มต้นบางอย่างเพื่อเชื่อมต่อกับคลัสเตอร์ระยะไกลและเริ่มต้น TPU หมายเหตุว่า tpu อาร์กิวเมนต์ TPUClusterResolver ที่อยู่พิเศษเฉพาะสำหรับ Colab ในกรณีที่คุณใช้งาน Google Compute Engine (GCE) คุณควรส่งในชื่อ CloudTPU ของคุณแทน

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
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.2:8470

INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.2: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: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')]

ตำแหน่งอุปกรณ์ด้วยตนเอง

หลังจากเริ่มต้น 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:localhost/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:XLA_CPU:0, XLA_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.experimental.TPUStrategy เพื่อขับเคลื่อน Cloud TPU และฝึกโมเดล Keras

กำหนดโมเดล Keras

ด้านล่างนี้คือคำจำกัดความของโมเดล MNIST โดยใช้ Keras ซึ่งไม่เปลี่ยนแปลงจากสิ่งที่คุณจะใช้กับ CPU หรือ 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)])

ใส่ชุดข้อมูล

การใช้tf.data.Dataset อย่างมีประสิทธิภาพ API ข้อมูลเป็นสิ่งสำคัญเมื่อใช้ Cloud TPU เนื่องจากเป็นไปไม่ได้ที่จะใช้ Cloud TPU เว้นแต่คุณจะป้อนข้อมูลได้เร็วพอ ดู คู่มือประสิทธิภาพการป้อนข้อมูลไปป์ไลน์ สำหรับรายละเอียดเกี่ยวกับประสิทธิภาพของชุดข้อมูล

สำหรับการทดลองทั้งหมดยกเว้นการทดลองที่ง่ายที่สุด (โดยใช้ tf.data.Dataset.from_tensor_slices หรือข้อมูลในกราฟอื่น ๆ ) คุณจะต้องจัดเก็บไฟล์ข้อมูลทั้งหมดที่อ่านโดยชุดข้อมูลในที่เก็บข้อมูล Google Cloud Storage (GCS)

สำหรับกรณีการใช้งานส่วนใหญ่ขอแนะนำให้แปลงข้อมูลของคุณเป็นรูปแบบ TFRecord และใช้ tf.data.TFRecordDataset เพื่ออ่าน ดู TFRecord และ tf ตัวอย่างบทช่วยสอน สำหรับรายละเอียดเกี่ยวกับวิธีการดำเนินการนี้ อย่างไรก็ตามนี่ไม่ใช่ข้อกำหนดที่เข้มงวดและคุณสามารถใช้โปรแกรมอ่านชุดข้อมูลอื่น ๆ ( FixedLengthRecordDataset หรือ TextLineDataset ) ได้หากต้องการ

ชุดข้อมูลขนาดเล็กสามารถโหลดลงในหน่วยความจำได้ทั้งหมดโดยใช้ tf.data.Dataset.cache

ไม่ว่าจะใช้รูปแบบข้อมูลใดขอแนะนำอย่างยิ่งให้คุณใช้ไฟล์ขนาดใหญ่โดยเรียงลำดับ 100MB นี่เป็นสิ่งสำคัญอย่างยิ่งในการตั้งค่าเครือข่ายนี้เนื่องจากค่าใช้จ่ายในการเปิดไฟล์สูงกว่ามาก

ที่นี่คุณควรใช้โมดูล 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

ฝึกโมเดลโดยใช้ Keras ระดับสูง API

คุณสามารถฝึกโมเดลได้ง่ายๆด้วย Keras fit / compile APIs ไม่มีอะไรที่เฉพาะเจาะจงของ TPU คุณจะต้องเขียนโค้ดเดียวกันด้านล่างหากคุณมี GPU หลายตัวและใช้ 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
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/data/ops/multi_device_iterator_ops.py:601: get_next_as_optional (from tensorflow.python.data.ops.iterator_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.data.Iterator.get_next_as_optional()` instead.

Warning:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/data/ops/multi_device_iterator_ops.py:601: get_next_as_optional (from tensorflow.python.data.ops.iterator_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.data.Iterator.get_next_as_optional()` instead.

  1/300 [..............................] - ETA: 20:25 - loss: 2.3062 - sparse_categorical_accuracy: 0.0500WARNING:tensorflow:Callbacks method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0015s vs `on_train_batch_end` time: 0.0237s). Check your callbacks.

Warning:tensorflow:Callbacks method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0015s vs `on_train_batch_end` time: 0.0237s). Check your callbacks.

300/300 [==============================] - ETA: 0s - loss: 0.1450 - sparse_categorical_accuracy: 0.9561WARNING:tensorflow:Callbacks method `on_test_batch_end` is slow compared to the batch time (batch time: 0.0013s vs `on_test_batch_end` time: 0.0097s). Check your callbacks.

Warning:tensorflow:Callbacks method `on_test_batch_end` is slow compared to the batch time (batch time: 0.0013s vs `on_test_batch_end` time: 0.0097s). Check your callbacks.

300/300 [==============================] - 15s 50ms/step - loss: 0.1450 - sparse_categorical_accuracy: 0.9561 - val_loss: 0.0456 - val_sparse_categorical_accuracy: 0.9852
Epoch 2/5
300/300 [==============================] - 8s 28ms/step - loss: 0.0357 - sparse_categorical_accuracy: 0.9890 - val_loss: 0.0372 - val_sparse_categorical_accuracy: 0.9884
Epoch 3/5
300/300 [==============================] - 9s 28ms/step - loss: 0.0193 - sparse_categorical_accuracy: 0.9940 - val_loss: 0.0557 - val_sparse_categorical_accuracy: 0.9835
Epoch 4/5
300/300 [==============================] - 9s 29ms/step - loss: 0.0141 - sparse_categorical_accuracy: 0.9954 - val_loss: 0.0405 - val_sparse_categorical_accuracy: 0.9883
Epoch 5/5
300/300 [==============================] - 9s 29ms/step - loss: 0.0092 - sparse_categorical_accuracy: 0.9967 - val_loss: 0.0428 - val_sparse_categorical_accuracy: 0.9887

<tensorflow.python.keras.callbacks.History at 0x7f480c793240>

เพื่อลดค่าใช้จ่ายหลามและเพิ่มประสิทธิภาพการทำงานของ 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)
Epoch 1/5
300/300 [==============================] - 16s 52ms/step - loss: 0.1337 - sparse_categorical_accuracy: 0.9583 - val_loss: 0.0521 - val_sparse_categorical_accuracy: 0.9840
Epoch 2/5
300/300 [==============================] - 5s 16ms/step - loss: 0.0331 - sparse_categorical_accuracy: 0.9898 - val_loss: 0.0360 - val_sparse_categorical_accuracy: 0.9884
Epoch 3/5
300/300 [==============================] - 5s 16ms/step - loss: 0.0188 - sparse_categorical_accuracy: 0.9939 - val_loss: 0.0405 - val_sparse_categorical_accuracy: 0.9887
Epoch 4/5
300/300 [==============================] - 5s 16ms/step - loss: 0.0117 - sparse_categorical_accuracy: 0.9961 - val_loss: 0.0808 - val_sparse_categorical_accuracy: 0.9820
Epoch 5/5
300/300 [==============================] - 5s 16ms/step - loss: 0.0119 - sparse_categorical_accuracy: 0.9962 - val_loss: 0.0488 - val_sparse_categorical_accuracy: 0.9862

<tensorflow.python.keras.callbacks.History at 0x7f47ac791240>

ฝึกโมเดลโดยใช้ลูปการฝึกที่กำหนดเอง

คุณยังสามารถสร้างและฝึกโมเดลของคุณโดยใช้ tf.function และ tf.distribute API ได้โดยตรง strategy.experimental_distribute_datasets_from_function API ใช้เพื่อแจกจ่ายชุดข้อมูลที่กำหนดฟังก์ชันชุดข้อมูล โปรดทราบว่าขนาดชุดงานที่ส่งไปยังชุดข้อมูลจะเป็นขนาดชุดงานจำลองแทนขนาดชุดงานส่วนกลางในกรณีนี้ หากต้องการเรียนรู้เพิ่มเติมโปรดดูการ ฝึกอบรมแบบกำหนดเองด้วย บทแนะนำ tf.distribute.Strategy

ขั้นแรกให้สร้างโมเดลชุดข้อมูลและฟังก์ชัน tf

# 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),))

จากนั้นเรียกใช้ลูปการฝึกอบรม

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.1335, accuracy: 95.8%
Epoch: 1/5
Current step: 600, training loss: 0.0344, accuracy: 98.93%
Epoch: 2/5
Current step: 900, training loss: 0.0196, accuracy: 99.35%
Epoch: 3/5
Current step: 1200, training loss: 0.0119, accuracy: 99.61%
Epoch: 4/5
Current step: 1500, training loss: 0.0102, accuracy: 99.65%

การปรับปรุงประสิทธิภาพหลายขั้นตอนภายใน tf.function

สามารถปรับปรุงประสิทธิภาพได้โดยการรันหลายขั้นตอนภายใน tf.function สิ่งนี้ทำได้โดยการตัดการเรียก strategy.run ด้วย tf.range ภายใน tf.function . AutoGraph จะแปลงเป็น 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.0087, accuracy: 99.72%

ขั้นตอนถัดไป