Я пытаюсь обучить простую модель MLP на Google Colab с использованием TPU.
Тем не менее, когда я пытаюсь конвертировать модель с
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
from keras.constraints import NonNeg
model = Sequential()
model.add(Dense(57,input_shape=(57,)))
model.add(Dense(60,kernel_constraint=NonNeg(),activation="relu"))
model.add(Dense(100,kernel_constraint=NonNeg(),activation="relu"))
model.add(Dense(50,kernel_constraint=NonNeg(),activation="relu"))
model.add(Dense(3,activation="linear"))
model.compile(optimizer='adam',loss='mse')
TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR']
tpu_model = tf.contrib.tpu.keras_to_tpu_model(
model,
strategy=tf.contrib.tpu.TPUDistributionStrategy(
tf.contrib.cluster_resolver.TPUClusterResolver(TPU_WORKER)))
выдает ошибку InvalidArgumentError: /job:localhost/replica:0/task:0/device:TPU_SYSTEM:0 unknown device.
полная ошибка здесь .
Это не имеет смысла для меня, потому что я скопировал код преобразования модели прямо из учебника Google. Когда я запускаю тестовый код:
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)
Я получаю ожидаемый результат:
TPU address is grpc://10.0.203.10:8470
TPU devices:
[_DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:CPU:0, CPU, -1, 1941375595625340814),
_DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 17179869184, 9079881000847066378),
_DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:0, TPU, 17179869184, 6922694346479333534),
_DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:1, TPU, 17179869184, 14324637633413341896),
_DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:2, TPU, 17179869184, 3528106575831937158),
_DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:3, TPU, 17179869184, 13852141601322651612),
_DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:4, TPU, 17179869184, 10344791506504172772),
_DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:5, TPU, 17179869184, 16666353711371098164),
_DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:6, TPU, 17179869184, 3428083526573573796),
_DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:7, TPU, 17179869184, 8632908473312514763),
_DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 17179869184, 9715206562754100387)]
Поскольку это работает, я не понимаю, почему преобразование модели не работает.