时间戳

首页 > 解决方案 > Google Cloud VM 仍然使用 CPU 而不是 TPU 来执行 Python/Tensorflow 脚本

python - Google Cloud VM 仍然使用 CPU 而不是 TPU 来执行 Python/Tensorflow 脚本

问题描述

我已经在 Google Cloud 上设置了一台 TPU 机器，并且我认为我已经正确完成了它，因为当我运行ctpu status它时会返回RUNNING.

但是，我有一个要运行的 Python 脚本，我希望它使用 TPU。根据终端输出的前几行，它仍在使用 CPU。输出是

[name: "/device:CPU:0"
device_type: "CPU"
memory_limit: 268435456
locality {
}
incarnation: 7191847218438877393
, name: "/device:XLA_CPU:0"
device_type: "XLA_CPU"
memory_limit: 17179869184
locality {
}
incarnation: 6058516396679200559
physical_device_desc: "device: XLA_CPU device"
]

我正在运行的命令是： python3 test.py 1 --tpu-name=$TPU_NAME

我已经运行export TPU_NAME=tpu_vm1并确认了echo

那么我可能做错了什么？我怎么能让脚本改用 TPU？

以防万一，除了我的test.py脚本之外，这里有一个删节：

#
# resnet time-to-accuracy-improvement tests
#

import os
from numpy.random import seed
seed(1)
import tensorflow as tf
tf.random.set_seed(2)
import numpy
import time

import tensorflow as tf
from tensorflow.keras.applications.resnet50 import ResNet50
import mycallbacks

from tensorflow.keras.applications.inception_v3 import preprocess_input
from tensorflow.keras.preprocessing.image import img_to_array
from tensorflow.keras.preprocessing.image import load_img
from tensorflow.keras.layers import Input
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.utils import plot_model

...

# display device type
from tensorflow.python.client import device_lib
print(device_lib.list_local_devices())

标签： pythontensorflowgoogle-cloud-platformtpugoogle-cloud-tpu

解决方案

因此，您列出了运行协调器的 GCE VM 上的本地 CPU 设备。TPU 本身由在 TPU 主机上运行的远程设备组成，而不是本地设备。

请查看此 Colab 笔记本。当你运行类似的东西时：

tpu = tf.distribute.cluster_resolver.TPUClusterResolver()  # TPU detection
tf.config.experimental_connect_to_cluster(tpu)
tf.tpu.experimental.initialize_tpu_system(tpu)
tpu_strategy = tf.distribute.experimental.TPUStrategy(tpu)

它会吐出TPUSystemMetadata，如：

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)
...

要将模型实际放置在 TPU 设备上，请确保使用此处描述的 TPU 策略。

推荐阅读