使用Google Colab和TPU训练物体检测模型

Question

我正在按照这里的指南来训练一个目标检测模型：https://towardsdatascience.com/creating-your-own-object-detector-ad69dda69c85。

在谷歌 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。我该如何做到这一点呢？

- migueljuan

---

你找到任何解决方案了吗？ - john doe

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.keras_to_tpu模块将其转换为TPU-model。

可以使用以下代码完成：

# 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链接和这个链接。

- user11530462