gpt4 book ai didi

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

转载 作者:行者123 更新时间:2023-12-04 08:26:11 28 4
gpt4 key购买 nike

我正在尝试使用 TPU 在 Colab 上训练和运行图像分类模型。没有pytorch。
我知道 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我认为这是解决方案,但我以前从未这样做过,我不知道从哪里开始。
有人可以就这个问题给我建议吗?

最佳答案

您必须使用 strategy.run() 调用此过程:

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

关于python - 使用分布式策略在 Colab TPU 上训练模型,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/65251058/

28 4 0
Copyright 2021 - 2024 cfsdn All Rights Reserved 蜀ICP备2022000587号
广告合作:1813099741@qq.com 6ren.com