"Unavailable: Socket closed" 在 TPU 上训练时出错

Question

在 TPU 上训练我的变压器模型时出现以下错误：

UnavailableError: 2 root error(s) found.
  (0) Unavailable: Socket closed
  (1) Invalid argument: Unable to find a context_id matching the specified one (13089686768223941123). Perhaps the worker was restarted, or the context was GC'd?

我的数据根据​​序列长度划分到桶中以获得最佳性能：

1. 长度小于等于8

2. 从 9 到 16

3. 从 17 到 24

我正在从随机桶中加载每批次。

当我第一次访问每个桶时 - tensorflow 内核创建一个新图并回溯模型。

第三次回溯时出现错误。所以如果我从任何两个桶训练我都没有错误。

Answer 1

据我了解 - 这是 tf 2.3 中的错误。

我切换到 2.2.0，错误消失了。