在Keras上使用
我们已经将 tf.distribute.Strategy 集成到 tf.keras 中。tf.keras 是一个构建和训练模型的高级API。通过集成到 tf.keras 后端, 用Keras训练框架写的程序可以无缝进行分布式训练。
您需要对代码中进行以下更改:
创建一个 tf.distribute.Strategy 实例
将Keras模型的创建和编译过程挪到strategy.scope中
支持各种类型的Keras模型:顺序模型、函数式模型和子类模型:
下面是一个非常简单的Keras模型示例:
mirrored_strategy = tf.distribute.MirroredStrategy()
with mirrored_strategy.scope():
model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])
model.compile(loss='mse', optimizer='sgd')
dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100).batch(10)
model.fit(dataset, epochs=2)
model.evaluate(dataset)
'''
Epoch 1/2
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
10/10 [==============================] - 0s 2ms/step - loss: 1.4031
Epoch 2/2
10/10 [==============================] - 0s 1ms/step - loss: 0.6202
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
10/10 [==============================] - 0s 1ms/step - loss: 0.3851
0.3851303160190582
'''
在自定义训练循环上使用
在高级 API上使用tf.distribute.Strategy 只需要修改几行代码. 再多花点功夫,你也可以在自定义训练循环上应用 tf.distribute.Strategy.
如果您需要比Estimator或Keras更大的灵活性和对训练循环更强的控制,则可以编写自定义训练循环。例如,使用GAN时,您可能希望每轮生成器和鉴别器训练不同的steps。同样,高级框架也不太适合强化学习。
为了支持自定义训练循环, 我们通过 tf.distribute.Strategy 类提供了一组核心方法。刚开始使用这些代码需要对代码进行较小的重组,但是一旦完成,只需更改Strategy实例就应该能够在GPU、TPU和多台计算机之间切换。
在这里,我们将展示一个简短的代码片段说明这个用例:使用与以前一样的Keras模型。
首先,需要在策略范围内创建模型和Optimizer。这样可以确保使用模型和Optimizer创建的任何变量都是镜像变量。
with mirrored_strategy.scope():
model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])
optimizer = tf.keras.optimizers.SGD()
#接下来, 我们要创建输入数据集并调用 tf.distribute.Strategy.experimental_distribute_dataset 将数据集按此策略分布:
dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(1000).batch(
global_batch_size)
dist_dataset = mirrored_strategy.experimental_distribute_dataset(dataset)
#然后定义训练的一个step,使用 tf.GradientTape 计算梯度,使用optimizer将梯度用于更新model变量。要分发这个训练步骤, 我们把它放到 step_fn 函数中,将次函数和从 dist_dataset 创建的数据集一起传递给 tf.distrbute.Strategy.run :
@tf.function
def train_step(dist_inputs):
def step_fn(inputs):
features, labels = inputs
with tf.GradientTape() as tape:
# training=True is only needed if there are layers with different
# behavior during training versus inference (e.g. Dropout).
logits = model(features, training=True)
cross_entropy = tf.nn.softmax_cross_entropy_with_logits(
logits=logits, labels=labels)
loss = tf.reduce_sum(cross_entropy) * (1.0 / global_batch_size)
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))
return cross_entropy
per_example_losses = mirrored_strategy.run(step_fn, args=(dist_inputs,))
mean_loss = mirrored_strategy.reduce(
tf.distribute.ReduceOp.MEAN, per_example_losses, axis=0)
return mean_loss
#定义好training_step之后, 我们就可以迭代 dist_dataset 循环进行训练
with mirrored_strategy.scope():
for inputs in dist_dataset:
print(train_step(inputs))