我想知道是否有办法确认我的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)