python - 确认 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 版本)，list_physical_devices('GPU') 也会返回 False，甚至如果您的系统物理上有 GPU。

“gpu一旦被TF识别就自动吗？”

是的。在 TF docs 之后引用:

Note: Use tf.config.experimental.list_physical_devices('GPU') to confirm that TensorFlow is using the GPU.

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

tf.debugging.set_log_device_placement(True)