tensorflow - 使用 TPU 在 Google Colab 上训练目标检测模型

Question

我正在按照此处的指南训练对象检测模型 https://towardsdatascience.com/creating-your-own-object-detector-ad69dda69c85

在 Google Colab 上，我能够执行以下操作并且它使用了 GPU

python train.py --train_dir=training/ --pipeline_config_path=training/ssd_mobilenet_v1_0.75_depth_quantized_300x300_coco14_sync.config

我现在想使用 TPU 进行训练，但这显然不是开箱即用的。运行 train.py 很慢并且似乎只使用 CPU。我怎样才能做到这一点？

Answer 1

在 Google Colab 中使用 TPU 时，我们应该使用下面提到的代码来检查 TPU 设备是否在环境中被正确识别:

import os
import pprint
import tensorflow as tf

if 'COLAB_TPU_ADDR' not in os.environ:
  print('ERROR: Not connected to a TPU runtime; please see the first cell in this notebook for instructions!')
else:
  tpu_address = 'grpc://' + os.environ['COLAB_TPU_ADDR']
  print ('TPU address is', tpu_address)

  with tf.Session(tpu_address) as session:
    devices = session.list_devices()

  print('TPU devices:')
  pprint.pprint(devices)

这应该会输出我们的 Colab 环境中可用的 8 个 TPU 设备的列表。

为了在 TPU 上运行 tf.keras 模型，我们必须使用 tf.contrib.tpu 将其转换为 TPU-model .keras_to_tpu 模块。

可以使用下面的代码来完成:

# This address identifies the TPU we'll use when configuring TensorFlow.
TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR']
tf.logging.set_verbosity(tf.logging.INFO)

resnet_model = tf.contrib.tpu.keras_to_tpu_model(
    resnet_model,
    strategy=tf.contrib.tpu.TPUDistributionStrategy(
        tf.contrib.cluster_resolver.TPUClusterResolver(TPU_WORKER)))

更多信息，请引用此Medium Link和 this Link .