为何TPU strategy.scope()会改变Keras model.fit()的输入数据类型？

原本train_X是np.ndarray类型（形状(16,128,128,3)），未使用TPU策略作用域时训练正常，但在tpu_strategy.scope()内定义模型后，train_X被自动转为BatchDataset，导致validation_split参数报错，示例如下：

# Case 1 正常训练
model = unet()
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])  
results = model.fit(train_X, train_Y, batch_size = 16, epochs = 100, validation_split=0.1, callbacks=callbacks)
>>> Starts training


# Case 2 触发报错
with tpu_strategy.scope():
   model = unet()
   model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

results = model.fit(train_X, train_Y, batch_size = 16, epochs = 100, validation_split=0.1)    
>>> `validation_split` argument is not supported when input `x` is a dataset or a dataset iterator. Received: x=<BatchDataset element_spec=(TensorSpec(shape=(16, 128, 128, 3), dtype=tf.uint8, name=None), TensorSpec(shape=(16, 128, 128, 1), dtype=tf.bool, name=None))>, validation_split=0.100000

请问该现象的原因是什么？

这是因为TPU策略会自动将numpy数组输入转换为TensorFlow Dataset对象，目的是适配TPU的分布式训练机制——TPU需要数据以Dataset形式实现高效的分布式分发、预处理和流水线加载，才能最大化利用TPU的计算性能。

在Case 1中，模型运行在CPU/GPU环境下，model.fit()可以直接处理numpy数组，并且支持validation_split参数自动划分验证集；而Case 2中，模型在TPU策略作用域内创建，Keras会自动触发输入转换逻辑，把传入的numpy数组封装成BatchDataset，但Dataset类型的输入并不支持validation_split参数，因此触发报错。

如果需要保留验证集划分逻辑，可以选择两种方式：

• 提前手动将numpy数组切分为训练集和验证集，分别传入model.fit()的x和validation_data参数
• 先将numpy数组转为Dataset对象，再使用Dataset的split方法手动划分训练/验证集

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