TPU 训练在某些指标上失败，在 CPU 上成功

Question

我正在尝试在一些图像上训练一个简单的 EfficientNet 样式模型。训练在 CPU 上运行良好，但是当我切换到使用 TPU 时，我收到以下错误：

(0) 无效参数：{{function_node __inference_train_function_38255}} then 和 else 分支的输出形状不匹配：(s64[1,

[[{{节点条件}}]]
[[TPUReplicate/_compile/_5430787790498024493/_4]]
[[tpu_compile_succeeded_assert/_6318656678166656164/_5/_289]]

(1) 无效参数：{{function_node __inference_train_function_38255}} then 和 else 分支的输出形状不匹配：(s64[1,

[[{{节点条件}}]]
[[TPUReplicate/_compile/_5430787790498024493/_4]]
[[tpu_compile_succeeded_assert/_6318656678166656164/_5/_225]]

仅当我使用特定指标Cohen's Kappa 时才会出现此错误。如果我删除这个指标，模型训练得很好。

我试图找出 CohensKappa 中的违规部分并将其缩小到 _update_confusion_matrix - 如果我超载此部分和 result，模型可以正常训练。

当我开始训练时，我看到以下日志消息：

TPU 具有动态形状的输入：[, , ]

这可能是相关的，但是考虑到当我忽略这个指标并且我仍然得到那个日志时模型训练得很好，这可能是一个红鲱鱼。

任何有关解决方案或如何调试的建议都会非常有帮助。切换到 Eager Execution 模式不是一种选择，因为它在 CPU 上运行良好。

Answer 1

请分享导致此错误的代码 sn-p。从它显示的内容来看，您似乎有张量维度不一致的问题（即(s64[1,<=4]) vs. (s64[<=4])）