Skip to main content
 首页 » 编程设计

python之无法在 Colab 中使用 TF 2.0 和 GPU 支持运行基本示例

2025年05月04日21unruledboy

较新版本的 TF (v2.0) 正在向 XLA 设备迁移,但这会导致 TF 丢失此类设备。 Colab 基本示例(可在此处找到)https://colab.research.google.com/notebooks/gpu.ipynb#scrollTo=v3fE7KmKRDsH ) 即使在环境选项选项卡中打开 GPU,它似乎也不起作用。问题仍然存在,模型正在 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

(无法为您提供数据集,它很大,抱歉......也为代码质量感到抱歉,它不应该被除我之外的任何人查看......)

我做错了什么吗?

请您参考如下方法:

问题是最近在Google Colab上安装了CUDA 10.1,最新稳定编译的TensorFlow 2不支持CUDA 10.1。解决方法是使用 Tensorflow 2 的夜间预览,因此请使用:

pip install tf-nightly-gpu