如何使用 Tensorflow 使用带有自定义损失的分布式训练？答案

【问题标题】：How to use distributed training with a custom loss using Tensorflow?如何使用 Tensorflow 使用带有自定义损失的分布式训练？
【发布时间】：2020-11-15 15:18:00
【问题描述】：

我有一个变压器模型，我想使用 Actor-Critic RL 进行训练，分布在 Google Cloud AI 平台上的多个工作人员中进行训练。我将我的数据按日期分解为单个文件并上传到 Cloud Storage。由于我使用的是 Actor-Critic RL，因此我有一个自定义损失函数来计算和应用梯度。我遇到的所有分布式训练示例都使用了model.fit，我将无法做到这一点。我找不到任何关于使用自定义损失的信息。

除了将训练分布在多台机器上之外，我还想知道如何在多个 CPU 内核之间正确分布训练。据我了解，model.fit 会正常处理这些事情。

这是自定义损失函数；现在它相当于批量大小 1 我相信：

def learn(self, state_value_starting: tf.Tensor, probabilities: tf.Tensor, state_new: tf.Tensor,
            reward: tf.Tensor, is_done: tf.Tensor):
    with tf.GradientTape() as tape:
        state_value_starting = tf.squeeze(state_value_starting)
        state_value_new, _ = self.call(state_new)
        state_value_new = tf.squeeze(state_value_new)

        action_probabilities = tfp.distributions.Categorical(probs=probabilities)
        log_probability = action_probabilities.log_prob(self._last_action)

        delta = reward + (self._discount_factor * state_value_new * (1 - int(is_done))) - state_value_starting
        actor_loss = -log_probability * delta
        critic_loss = delta ** 2
        total_loss = actor_loss + critic_loss

    gradient = tape.gradient(total_loss, self.trainable_variables)
    self.optimizer.apply_gradients(zip(gradient, self.trainable_variables))

【问题讨论】：

标签： python tensorflow keras distributed-computing

【解决方案1】：

Tensorflow Model 提供了一个实践的解决方案，在model_lib_v2.py 中定义。

见train_loop函数，自定义训练循环构造利用

strategy = tf.compat.v2.distribute.get_strategy() #L501
with strategy.scope():
    training step ...

以及函数eager_train_step中的自定义损失。

【讨论】：