google-compute-engine - TPU 训练在某些指标上失败，在 CPU 上成功

Question

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

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

[[{{node cond}}]]
[[TPUReplicate/_compile/_5430787790498024493/_4]]
[[tpu_compile_succeeded_assert/_6318656678166656164/_5/_289]]

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

[[{{node cond}}]]
[[TPUReplicate/_compile/_5430787790498024493/_4]]
[[tpu_compile_succeeded_assert/_6318656678166656164/_5/_225]]

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

我试图找出 CohensKappa 中的违规部分并将其缩小到_update_confusion_matrix- 如果我超载这个和result，模型训练得很好。

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

TPU 具有动态形状的输入： [<tf.Tensor 'Const:0' shape=() dtype=int32>, <tf.Tensor 'cond_8/Identity:0' shape=(None, 456, 456, 3) dtype= float32>, <tf.Tensor 'cond_8/Identity_1:0' shape=(None,) dtype=int64>]

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

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

标签： google-compute-enginegoogle-cloud-tpu

Answer 1

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