确认TF2在训练时使用我的GPU

Question

我想知道是否有办法确认我的TF模型在GPU上训练，因为我按照TF教程的建议将训练数据存储在GPU上。以下是一个简短的代码示例：

import tensorflow as tf

print('Num GPUs Available:', len(tf.config.experimental.list_physical_devices('GPU')))

# load data on GPU
with tf.device('/GPU:0'):
    mnist = tf.keras.datasets.mnist
    (x_train, y_train), (x_test, y_test) = mnist.load_data()
    x_train, x_test = x_train / 255.0, x_test / 255.0
# define, compile and train the model
model = tf.keras.models.Sequential([tf.keras.layers.Dense(1)])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['acc'])
model.fit(x_train, y_train, batch_size=32, epochs=5)

- Gepeto97

---

你可以随时运行nvidia-smi或类似的命令行工具来查看。 - chris

Answer 1

在Tensorflow 2.x中，有几种检查GPU的方法。如果GPU可用，则模型将在其上运行（除非它正在被另一个锁定它的TF实例占用）。放置也会显示在日志文件中，并且可以通过例如nvidia-smi进行确认。

在下面的代码中，我将假设tensorflow已经按照惯例和您的代码导入为tf。

要检查哪些设备可用，请运行：

tf.config.experimental.list_physical_devices()

以下是我的输出：

[物理设备（名称= '/physical_device:CPU:0'，设备类型='CPU'）， 物理设备（名称='/physical_device:XLA_CPU:0'，设备类型='XLA_CPU'）， 物理设备（名称='/physical_device:GPU:0'，设备类型='GPU'）， 物理设备（名称='/physical_device:XLA_GPU:0'，设备类型='XLA_GPU'）]

为了检查系统中是否有任何GPU:

is_gpu = len(tf.config.experimental.list_physical_devices('GPU')) > 0

从Tensorflow 2.1开始，此功能已从实验性迁移，您可以以相同的方式使用tf.config.list_physical_devices()。

is_gpu = len(tf.config.list_physical_devices('GPU')) > 0 

在某个时间点，实验部分将被弃用。

最后但并非最不重要的是，如果你的tensorflow没有使用CUDA编译（即非GPU版本），list_physical_devices('GPU')也会返回False，即使你的系统实际上有一个GPU。

“一旦TF识别了GPU，它就是自动的吗？”

是的。引用自TF文档：

注意：使用tf.config.experimental.list_physical_devices('GPU')来确认TensorFlow正在使用GPU。

如果被识别，它将在训练期间被使用。如果您想确保，您可以要求更明确的日志记录：

tf.debugging.set_log_device_placement(True)

- Lukasz Tracewski

---

谢谢，但我如何确保训练正在使用GPU？如果TF自动识别GPU，那么它是自动的吗？ - Gepeto97

答案是“是”。稍微详细些的解释在编辑后的文本中，最后一节。 - Lukasz Tracewski

Answer 2

有一种更简单的方法来实现这个：

import tensorflow as tf
if tf.test.gpu_device_name():
    print('Default GPU Device: {}'.format(tf.test.gpu_device_name()))
else:
    print(""Please install GPU version of TF"")

（或者）

sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))

(或者)

TF中出现了一些有用的函数：

判断GPU是否可用

tf.test.is_gpu_available()

返回GPU设备的名称。

tf.test.gpu_device_name()  

- bsquare

---

从 https://www.tensorflow.org/api_docs/python/tf/test/is_gpu_available : tf.test.is_gpu_available : 警告：此函数已被弃用。它将在未来的版本中被删除。更新说明：请改用 tf.config.list_physical_devices('GPU')。另外，tf.Session 仅适用于 TF1。 - Lukasz Tracewski

Answer 3

现在是8月23日，使用以下代码： tf.config.list_physical_devices('GPU') 来源于官方的TensorFlow安装页面。 tensorflow.org/install/pip

- Rohit Dhankar