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