使用分布式策略在 Colab TPU 上训练模型

Question

我正在尝试使用 TPU 在 Colab 上训练和运行图像分类模型。没有火炬。

我知道 TPU 仅适用于 GCS 存储桶中的文件，因此我从存储桶中加载数据集，并且我还评论了检查点和日志记录功能，以免出现此类错误。我只想看到它在 TPU 上没有错误地训练。

在 CPU 和 GPU 上代码有效，但是当我在创建模型之前使用 with strategy.scope(): 时出现问题。这是在我训练模型时给我带来问题的函数：

def train_step(self, images, labels):
    with tf.GradientTape() as tape:
        predictionProbs = self(images, training=True)
        loss = self.loss_fn(labels, predictionProbs)

    grads = tape.gradient(loss, self.trainable_weights)

    predictionLabels = tf.squeeze(tf.cast(predictionProbs > PROB_THRESHOLD_POSITIVE, tf.float32), axis=1)
    acc = tf.reduce_mean(tf.cast(predictionLabels == labels, tf.float32))

    self.optimizer.apply_gradients(zip(grads, self.trainable_weights)) # here is the problem

    return loss, acc

这是我遇到的错误：

RuntimeError: `apply_gradients() cannot be called in cross-replica context. Use `tf.distribute.Strategy.run` to enter replica context.

我看过 https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy，我认为这是解决方案，但我以前从未这样做过，我不知道从哪里开始。

有人可以就这个问题给我一些建议吗？

Answer 1

您必须使用 strategy.run():

调用此过程

strategy.run(train_step, args=(images, labels))

使用分布式策略在 Colab TPU 上训练模型

Train model on Colab TPU with distributed strategy

python

distributed-computing

tensorflow

tpu

tensorflow2.0