在TPU上进行训练时，验证损失变得微不足道，但在GPU上完全可以

如何解决在TPU上进行训练时，验证损失变得微不足道，但在GPU上完全可以

我在使用GPU和TPU的Google colab训练中遇到了一个相当奇怪的问题，我使用了自定义损失，并且在GPU上的tfrecord数据集还可以，但是如果我切换到TPU，则会损失nan作为验证损失。没有其他特定错误。此外，较早的验证tfrecord数据集在TPU上也可以正常工作。这使我认为可能存在一些特定于数据的东西。这也出现在model.evaluate（...）上，因为它在验证集上。

有人知道如何最好地用TPU调试吗？可根据要求提供更多详细信息。

解决方法

我的问题可能与https://github.com/tensorflow/tensorflow/issues/41635有关（尽管即使对于非自定义损失函数也可以看到）。就我而言，我看不到使用开箱即用损失功能，但是当我使用自定义损失时会遇到问题。自定义丢失似乎并不是主要原因，因为它对于任何数据集都可以在CPU和GPU下工作。

无论如何，我遵循问题​​的提示，并删除最后一批（其大小小于batch_size），并且不再显示NaN。虽然这可以解决问题，但对于根本原因，我仍然没有明确的答案。