无法在Colab中运行TF2.0和GPU支持的基本示例

Question

一个较新版本的TF (v2.0)正在向XLA设备迁移，但是这会导致TF失去这类设备。即使在环境选项卡中打开GPU，基本示例Colab (可以在这里找到https://colab.research.google.com/notebooks/gpu.ipynb#scrollTo=v3fE7KmKRDsH)似乎也无法工作。问题仍然存在，模型要么在CPU上运行，要么在通过使用with tf.device('/device:XLA_GPU:0'):显式指定时运行。

with tf.device('/device:GPU:0'): # <=== Device not found
  ...

with tf.device('/device:XLA_GPU:0'): # <=== This part runs fine but in attempt to run inference or training fails with exception (see below)
  base_model = tf.keras.applications.MobileNetV2(input_shape=IMG_SHAPE,
                                           include_top=False, 
                                           weights='imagenet') 

  maxpool_layer = tf.keras.layers.GlobalMaxPooling2D()
  prediction_layer = tf.keras.layers.Dense(1, activation='sigmoid')
  model = tf.keras.Sequential([
      base_model,
      maxpool_layer,
      prediction_layer
  ])

  model.compile(optimizer=tf.keras.optimizers.Adam(lr=learning_rate),
                loss='binary_crossentropy',
                metrics=['accuracy'])

  model.summary()

但是，除例外情况外，以下代码行将失败：

with tf.device('/device:XLA_GPU:0'):
  history = model.fit(train_data.repeat(),
                      epochs=num_epochs,
                      steps_per_epoch = steps_per_epoch)

抱歉没有记录异常，因为它需要大约一个小时的时间来再次拆分我的colab (这里已经5点了)，然而问题似乎很多，似乎没有办法使用Colab提供的GPU和TF 2.0。

这是我的colab的链接，它是一个简单的拇指向上检测器

https://colab.research.google.com/drive/1ycfeET5DeqpI3oaY1gt5PpalqPNmqECg

(无法提供数据集，很大，很抱歉.对代码的质量也很抱歉，它不应该被任何人看到，除了我.

我做错了什么吗？

Answer 1

问题是CUDA10.1是最近在Google安装的，而最新的TensorFlow 2稳定编译并不支持CUDA10.1。解决方法是使用Tensorflow 2的夜间预览，所以请使用：

pip install tf-nightly-gpu

Answer 2

我找到了答案，这里。这是来自Tensorflow的一个bug。

只需在笔记本顶部插入这一行并重新启动内核即可。

%tensorflow_version 2.x