seq2seq翻译模型CPU/GPU运行正常，TPU训练报ValueError: None values not supported

TPU运行时的分布式策略和CPU/GPU不同，触发报错的核心原因有两个：一是Python原生生成器无法被TPUStrategy正确序列化传输到TPU节点，二是输入数据的shape或值存在动态不确定性（如缺省的None值、动态batch维度），可按以下步骤修改：

• 替换输入数据源格式：弃用已废弃的fit_generator接口，将Python生成器转换为tf.data.Dataset格式，TPU要求每个输入batch的shape完全固定，必须配置drop_remainder=True丢弃不足batch大小的尾批次。
  参考代码：

# 定义数据集输出的固定shape和类型
output_signature = (
    (tf.TensorSpec(shape=(max_length_src,), dtype=tf.float32),
     tf.TensorSpec(shape=(max_length_tar,), dtype=tf.float32)),
    tf.TensorSpec(shape=(max_length_tar, num_decoder_tokens), dtype=tf.float32)
)

# 构造训练、验证数据集
train_dataset = tf.data.Dataset.from_generator(
    lambda: generate_batch(X_train, y_train, batch_size=batch_size),
    output_signature=output_signature
).batch(batch_size, drop_remainder=True).prefetch(tf.data.AUTOTUNE)

val_dataset = tf.data.Dataset.from_generator(
    lambda: generate_batch(X_test, y_test, batch_size=batch_size),
    output_signature=output_signature
).batch(batch_size, drop_remainder=True).prefetch(tf.data.AUTOTUNE)

# 修改训练调用
model.fit(train_dataset,
          steps_per_epoch = train_samples // batch_size,
          epochs = epochs,
          validation_data = val_dataset,
          validation_steps = val_samples // batch_size)

• 修复生成器的空值风险：在词索引查找时添加未登录词兜底逻辑，避免词表缺失对应词时返回空值触发报错。首先在构造input_token_index和target_token_index时加入<UNK>特殊标记，再修改生成器逻辑：

for t, word in enumerate(input_text.split()):
    encoder_input_data[i, t] = input_token_index.get(word, input_token_index['<UNK>'])
for t, word in enumerate(target_text.split()):
    target_word_idx = target_token_index.get(word, target_token_index['<UNK>'])
    if t < len(target_text.split())-1:
        decoder_input_data[i, t] = target_word_idx
    if t > 0:
        decoder_target_data[i, t - 1, target_word_idx] = 1.

• 兼容版本适配：TensorFlow 2.6的TPUStrategy存在多个已知的自定义输入兼容bug，如修改后仍报错可降级到TensorFlow 2.3版本，符合不使用1.x版本的要求。

内容的提问来源于stack exchange，提问作者Adarsh Wase