在 TensorFlow 里,多 GPU 编程可借助数据并行的方式实现,也就是把数据划分到不同的 GPU 上并行处理。下面是一个简单的多 GPU 编程示例:
import tensorflow as tf
from tensorflow.keras.datasets import mnist
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Flatten
from tensorflow.keras.optimizers import Adam
# 加载数据集
(x_train, y_train), (x_test, y_test) = mnist.load_data()
# 数据预处理
x_train = x_train / 255.0
x_test = x_test / 255.0
# 定义模型
def create_model():
model = Sequential([
Flatten(input_shape=(28, 28)),
Dense(128, activation='relu'),
Dense(10, activation='softmax')
])
model.compile(optimizer=Adam(),
loss='sparse_categorical_crossentropy',
metrics=['accuracy'])
return model
# 检测可用的 GPU
strategy = tf.distribute.MirroredStrategy()
print('Number of devices: {}'.format(strategy.num_replicas_in_sync))
# 在策略范围内创建模型
with strategy.scope():
model = create_model()
# 训练模型
model.fit(x_train, y_train, epochs=5, batch_size=64 * strategy.num_replicas_in_sync)
# 评估模型
test_loss, test_acc = model.evaluate(x_test, y_test)
print(f"Test accuracy: {test_acc}")
代码解释:
- 加载数据集:借助
mnist.load_data()
加载 MNIST 数据集。 - 数据预处理:把像素值归一化到 0 到 1 的范围。
- 定义模型:构建一个简单的全连接神经网络模型。
- 检测可用的 GPU:使用
tf.distribute.MirroredStrategy()
来检测可用的 GPU 设备。 - 在策略范围内创建模型:使用
strategy.scope()
确保模型和优化器在多 GPU 环境下正确运行。 - 训练模型:调用
model.fit()
方法来训练模型。 - 评估模型:使用
model.evaluate()
方法评估模型在测试集上的性能。
这个示例展示了怎样在 TensorFlow 里使用 MirroredStrategy
进行多 GPU 训练。MirroredStrategy
会在每个 GPU 上复制一份模型,并且同步更新模型参数。