小编Gep*_*o97的帖子

我想知道在我按照 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)