【tensorflow2.0】49.Keras和自定义分布式训练

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