如何使用 Tensorflow 进行自定义损失的分布式训练？

Question

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

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

这是自定义损失 function； 我相信现在它相当于 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))

Answer 1

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

见 function train_loop ，构建自定义训练循环利用

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

自定义损失在 function eager_train_step 。