Я настроил машину TPU в Google Cloud и думаю, что все сделал правильно, потому что при запуске ctpu status возвращается RUNNING.
Однако у меня есть скрипт Python, который я пытаюсь запустить и хочу использовать ТПУ. Тем не менее, он все еще использует 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())