Kaggle TPU报错failed to connect to all addresses问题求助

错误原因说明 #

你遇到的两类报错存在关联：

• 首次出现的CancelledError属于输入数据不符合TPU分布式训练要求导致的训练中断
• 后续的UnavailableError是首次训练报错后TPU集群连接被异常释放，无法再次复用现有会话连接导致的

具体修复步骤 #

• 输入数据适配TPU要求 #

你当前的代码应该是直接将numpy格式的数据集传入model.fit，没有做TPU兼容的数据集封装和分发，这是报错的核心原因。需要将训练、验证数据集都封装为tf.data.Dataset格式，并通过TPU策略做分布式分发：

# 假设x_train、y_train为预处理后的训练数据，x_val、y_val为验证数据
BATCH_SIZE_PER_REPLICA = 128
GLOBAL_BATCH_SIZE = BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync

# 封装训练数据集
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))\
                               .shuffle(10000)\
                               .batch(GLOBAL_BATCH_SIZE)\
                               .prefetch(tf.data.AUTOTUNE)
# 封装验证数据集
val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val))\
                             .batch(GLOBAL_BATCH_SIZE)\
                             .prefetch(tf.data.AUTOTUNE)

# 将数据集分发到TPU各个核心
train_dist_dataset = strategy.experimental_distribute_dataset(train_dataset)
val_dist_dataset = strategy.experimental_distribute_dataset(val_dataset)

训练时直接传入分发后的数据集即可：

Model.fit(train_dist_dataset, validation_data=val_dist_dataset, epochs=你的训练轮次)

• 修复TPU连接异常 #

首次报错后出现的连接错误不需要修改业务代码，直接重启Kaggle Notebook的运行会话，重新按顺序执行所有代码即可，不要在报错后重复执行TPU初始化、训练相关代码。

• 代码优化（非强制，可避免潜在问题） #

你当前将模型定义、编译拆分为两个独立的strategy.scope()块，可以合并为一个，保证所有变量都在TPU策略作用域内初始化：

with strategy.scope():
    Model = Sequential([
        InputLayer((28, 28, 1)),
        Dropout(0.1),
        Conv2D(128, 3, use_bias=False),
        LeakyReLU(0.05),
        BatchNormalization(),
        MaxPooling2D(2, 2),
        Conv2D(64, 3, use_bias=False),
        LeakyReLU(0.05),
        BatchNormalization(),
        MaxPooling2D(2, 2),
        Flatten(),
        Dense(128, use_bias=False),
        LeakyReLU(0.05),
        BatchNormalization(),
        Dense(10, activation='softmax')
    ])

    Model.compile(optimizer='adam',
                  loss='categorical_crossentropy', metrics='accuracy') 

• 注意点 #

TPU训练的全局batch size建议设置为8的倍数（Kaggle TPU默认包含8个计算核心），可以最大化TPU硬件利用率。

内容的提问来源于stack exchange，提问作者Dacian Peng