Colab TPU训练模型时出现UnimplementedError动态维度传播错误如何解决？

这个报错的核心是TPU不支持携带动态维度的张量执行跨副本求和（CrossReplicaSum）操作，你没有感知到动态维度是因为触发场景比较隐蔽，具体原因分为三类：

• 数据集输出形状不固定：你已经确认数据集可以正常拉取数据，但大概率是数据预处理时没有强制固定图像尺寸，或是生成批次时没有丢弃样本数不足BATCH_SIZE的最后一个批次，产生了运行时才能确定的动态维度（报错日志里的<=3就是动态维度的标记）。TPU要求所有输入张量的形状在编译阶段必须完全确定，GPU对动态形状的兼容性更强所以不会报错。
• steps_per_execution参数适配问题：你在模型编译时设置了steps_per_execution = 32，该参数在TPU上运行时要求总训练步数必须是该值的整数倍，如果len(df_train) // config.BATCH_SIZE不能被32整除，也会触发动态形状判定。
• 第三方EfficientNet实现的隐含动态操作：你使用的第三方efn库的EfficientNetB0，部分padding、下采样逻辑是动态计算的，GPU可以兼容但TPU不支持这类动态操作。

你可以按顺序尝试以下修复方案：

• 强制数据集输出静态形状：生成批次时开启丢弃不完整批次的配置，同时给张量添加固定形状约束，示例代码如下：

# 生成批次时丢弃不足batch size的最后一批
train_dataset = train_dataset.batch(config.BATCH_SIZE, drop_remainder=True)
valid_dataset = valid_dataset.batch(config.BATCH_SIZE, drop_remainder=True)

# 强制固定张量形状
def fix_shape(image, label):
    image.set_shape([config.BATCH_SIZE, config.IMG_SIZE, config.IMG_SIZE, 3])
    label.set_shape([config.BATCH_SIZE])
    return image, label

train_dataset = train_dataset.map(fix_shape)
valid_dataset = valid_dataset.map(fix_shape)

• 调整steps_per_execution配置：要么把该参数值改为1，要么调整steps_per_epoch的取值，确保其可以被steps_per_execution的取值整除。
• 替换EfficientNet实现：改用tensorflow.keras.applications内置的EfficientNetB0实现，内置版本已经完成了TPU静态形状适配，不会产生动态维度。

内容的提问来源于stack exchange，提问作者Abhishek Prajapat