python - 当我尝试转换 Keras MLP 时，为什么 Google Colab 会给出 "unknown device"错误？

Question

我正在尝试使用 TPU 在 google colab 上训练一个简单的 MLP 模型。 但是，当我尝试使用

转换模型时

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未知设备。 full error here .

这对我来说没有意义，因为我直接从谷歌教程复制了模型转换代码。当我运行测试代码时:

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

既然这样有效，我不明白为什么模型转换不起作用。

Answer 1

如果启用了急切执行，您可能希望将其关闭。
对我来说，在代码开头没有 tf.enable_eager_execution() 似乎进展顺利。