详解TensorFlow报”AbortedError: Multiple devices detected for operation “的原因以及解决办法

问题描述

使用 TensorFlow 运行代码时，可能会遇到以下报错：

AbortedError (see above for traceback): Multiple devices detected for operation XXX

报错提示含义

这个报错提示表示：TensorFlow 在操作 XXX 上检测到了多个设备。

可能原因

造成这个问题的原因可能有：

使用了多个 GPU

当你的电脑上有多个 GPU 时，由于 TensorFlow 默认会尽可能地使用所有 GPU 完成计算任务，因此可能会使用多个 GPU。

使用了 TensorFlow 的集群模式

当使用 TensorFlow 的集群模式时，也会在多个计算设备上执行计算任务。

程序中指定了多个设备

在不同操作中指定了不同的设备也可能导致这个问题。

解决办法

遇到这个问题时，可以使用以下方法进行解决：

确认操作所使用的设备

可以使用 tf.get_default_graph().as_graph_def() 函数打印出默认计算图中的所有操作及其设备信息，查看是否有某个操作使用了多个设备。

修改代码让某些操作只使用一个设备

可以使用 tf.device() 函数来指定操作所在的设备。例如，将所有操作都放在 CPU 上执行：

with tf.device('/cpu:0'):
    # 执行一些 TensorFlow 操作

修改代码只使用一个 GPU

在进行训练任务时，如果只使用一个 GPU 能够满足需求，则可以使用以下代码将 TensorFlow 强制限定在使用一个 GPU 上：

import os
os.environ["CUDA_VISIBLE_DEVICES"]="0"  # 只使用第 0 个 GPU

其中，数字 0 表示第一个 GPU，数字 1 表示第二个 GPU，以此类推。

修改代码关闭 TensorFlow 的集群模式

如果你不需要在分布式环境下使用 TensorFlow，则可以在代码中设置以下命令关闭集群模式：

tf.config.set_cluster_def(None)

指定环境变量关闭 TensorFlow 的多 GPU 功能

可以使用以下命令指定环境变量 CUDA_VISIBLE_DEVICES，将 TensorFlow 强制限制只使用一个 GPU：

CUDA_VISIBLE_DEVICES=0 python my_program.py

其中，数字 0 表示只使用第一个 GPU。

参考资料

• TensorFlow 官方文档：https://www.tensorflow.org/api_docs/python/tf/device
• TensorFlow 官方用户指南：https://www.tensorflow.org/guide/gpu#using_multiple_gpus
• TensorFlow 官方 GitHub Issue：https://github.com/tensorflow/tensorflow/issues/4192