在训练时确认 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)

Answer 1

在 Tensorflow 2.x 中有几种检查 GPU 的方法。本质上，如果 GPU 可用，那么模型将在其上运行（除非它正忙于例如锁定它的另一个 TF 实例）。该位置也将在日志文件中看到，并且可以使用例如nvidia-smi.

在下面的代码中，我将假设tensorflow导入为tf（根据约定和您的代码）。

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

tf.config.experimental.list_physical_devices()

这是我的输出：

[PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU'), PhysicalDevice(name='/physical_device:XLA_CPU:0', device_type='XLA_CPU'), PhysicalDevice(name='/physical_device:GPU :0', device_type='GPU'), PhysicalDevice(name='/physical_device:XLA_GPU:0', device_type='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 版本），即使您的系统物理上有一个 GPU ，list_physical_devices('GPU')也会返回False。

“一旦gpu被TF识别，它是自动的吗？”

是的。在TF 文档之后引用：

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

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

tf.debugging.set_log_device_placement(True)