在 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}")

代码解释:

  1. 加载数据集:借助 mnist.load_data() 加载 MNIST 数据集。
  2. 数据预处理:把像素值归一化到 0 到 1 的范围。
  3. 定义模型:构建一个简单的全连接神经网络模型。
  4. 检测可用的 GPU:使用 tf.distribute.MirroredStrategy() 来检测可用的 GPU 设备。
  5. 在策略范围内创建模型:使用 strategy.scope() 确保模型和优化器在多 GPU 环境下正确运行。
  6. 训练模型:调用 model.fit() 方法来训练模型。
  7. 评估模型:使用 model.evaluate() 方法评估模型在测试集上的性能。

这个示例展示了怎样在 TensorFlow 里使用 MirroredStrategy 进行多 GPU 训练。MirroredStrategy 会在每个 GPU 上复制一份模型,并且同步更新模型参数。