Colab TPU错误InvalidArgumentError:无法分配设备进行操作

时间:2019-10-26 20:30:11

标签: tensorflow keras google-colaboratory tpu google-cloud-tpu

在使用TPU的google colab中,我遇到以下错误

InvalidArgumentError:无法为操作Adam / iterations / IsInitialized / VarIsInitializedOp分配设备:{{node Adam / iterations / IsInitialized / VarIsInitializedOp}}已明确分配给/ job:worker / replica:0 / task:0 / device: TPU:0,但可用的设备是[/ job:localhost /副本:0 / task:0 / device:CPU:0,/ job:localhost /副本:0 / task:0 / device:XLA_CPU:0]。确保设备规格引用的是有效设备。      [[Adam / iterations / IsInitialized / VarIsInitializedOp]]

TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR']
        resolver = tf.contrib.cluster_resolver.TPUClusterResolver(TPU_WORKER)
        tf.contrib.distribute.initialize_tpu_system(resolver)
        strategy = tf.contrib.distribute.TPUStrategy(resolver)
        with strategy.scope():
            # Setup the model inputs / outputs
            model = Model(inputs=[inputs_img, inputs_mask], outputs=outputs)
            # Compile the model
            model.compile(
                optimizer = Adam(lr=lr),
                loss=self.loss_total(inputs_mask)
            )

2 个答案:

答案 0 :(得分:0)

keras.modelskeras.layers移到tf.keras.modelstf.keras.layers为我解决了这个问题。 还要确保您没有使用本地存储。请改用GCS存储桶。

答案 1 :(得分:0)

这里是参考TPU合作伙伴,请确保将加速器更改为TPU: 运行时->更改运行时间-> TPU

https://colab.research.google.com/notebooks/tpu.ipynb