在 TPU 上训练时验证损失变为 nan 但在 GPU 上完全没问题答案

【问题标题】：Validation loss become nan while training on TPU but perfectly ok on GPU在 TPU 上训练时验证损失变为 nan 但在 GPU 上完全没问题
【发布时间】：2021-01-12 17:06:19
【问题描述】：

我在使用 GPU 和 TPU 进行 google colab 训练时遇到了一个相当奇怪的问题，我使用了自定义损失，并且在 GPU 上的 tfrecord 数据集很好，但如果我切换到 TPU，则将 nan 作为验证损失。没有其他具体错误。此外，较旧的验证 tfrecord 数据集在 TPU 上运行良好。这让我觉得可能有一些特定于数据的东西。这也出现在 model.evaluate(...) 上，因为它在验证集上。

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

【问题讨论】：

我会尝试以下步骤：1) 使用随机输入/验证调用模型，2) 使用随机输入/验证拟合模型，3) 拟合简单损失函数，4) 简化自定义损失函数，直到它有效。

标签： google-colaboratory tensorflow2.0 nan tfrecord tpu

【解决方案1】：

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

无论如何，我按照问题的提示删除了最后一批（它的大小小于 batch_size）并且不再看到 NaN。虽然这解决了问题，但我对根本原因仍然没有明确的答案。

【讨论】：