在 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，错误消失了。