在Keras中,可以使用tf.config.experimental.set_visible_devices
函数来指定可见的GPU。下面是一个示例代码:
import tensorflow as tf
from tensorflow import keras
# 禁用GPU
tf.config.experimental.set_visible_devices([], 'GPU')
# 构建模型
model = keras.Sequential([
keras.layers.Dense(64, activation='relu'),
keras.layers.Dense(10, activation='softmax')
])
# 编译模型
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
# 加载数据并训练模型
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train = x_train.reshape(-1, 784) / 255.0
x_test = x_test.reshape(-1, 784) / 255.0
y_train = keras.utils.to_categorical(y_train)
y_test = keras.utils.to_categorical(y_test)
model.fit(x_train, y_train, batch_size=64, epochs=10, validation_data=(x_test, y_test))
在上面的代码中,tf.config.experimental.set_visible_devices([], 'GPU')
将禁用所有可见的GPU。这样,即使在计算机上有多个GPU,Keras也只会使用CPU来训练模型。