如何在TensorFlow函数式API中拆分Dataset元组为两个独立层

要在函数式API内部从返回(training_data, label)的BatchDataset中独立获取两个元素，核心思路是用单个输入层接收整个元组输入，再通过Lambda层或直接张量索引在模型内部拆分出训练数据和标签。以下是具体实现步骤：

1. 明确数据集的元素结构 #

首先确认数据集返回的元组各元素的形状与类型，这是定义输入层的依据：

import tensorflow as tf
import numpy as np

# 超参数定义
batch_size = 32
sequence_length = 10

# 构建玩具数据集
time = np.arange(1000)
data = np.random.randn(1000)
label = np.random.randn(1000)

training_data = np.zeros(shape=(time.size,2))
training_data[:,0] = time
training_data[:,1] = data

dataset_train = tf.keras.utils.timeseries_dataset_from_array(
      data = training_data,
      targets = label,
      batch_size = batch_size, 
      sequence_length = sequence_length,
      sequence_stride = 1,
  )

# 查看数据集元素规格
print(dataset_train.element_spec)
# 输出示例：(TensorSpec(shape=(None, 10, 2), dtype=tf.float64, name=None), TensorSpec(shape=(None,), dtype=tf.float64, name=None))

2. 构建支持元组输入的函数式模型 #

使用tf.keras.Input的type_spec参数直接匹配数据集的元组结构，再通过Lambda层拆分出训练数据和标签：

# 基于数据集的元素规格定义输入层
input_tuple = tf.keras.Input(type_spec=dataset_train.element_spec)

# 拆分元组，得到训练数据和标签
training_data = tf.keras.layers.Lambda(lambda x: x[0])(input_tuple)
label_tensor = tf.keras.layers.Lambda(lambda x: x[1])(input_tuple)

# 对训练数据执行时序模型操作（示例）
x = tf.keras.layers.LSTM(64, dtype=tf.float32)(training_data)
predictions = tf.keras.layers.Dense(1)(x)

# 示例：将标签也作为模型输出（根据需求调整输出结构）
model = tf.keras.Model(inputs=input_tuple, outputs=[predictions, label_tensor])

# 编译模型：根据输出结构定义损失函数，无需计算损失的输出可设为None
model.compile(
    optimizer=tf.keras.optimizers.Adam(),
    loss=['mse', None]  # 仅对预测结果计算MSE损失
)

3. 训练与验证 #

直接传入数据集训练即可，模型会自动识别元组输入并拆分：

# 训练模型
model.fit(dataset_train, epochs=5)

# 预测示例：传入数据集的单个批次
for batch in dataset_train.take(1):
    pred, lbl = model.predict(batch)
    print(f"预测结果形状：{pred.shape}，标签形状：{lbl.shape}")

对之前尝试方案的问题解析 #

• 双输入层方案：model.fit会将Dataset返回的元组第一个元素视为模型输入，第二个视为训练目标，而非同时传入两个输入层，因此出现输入数量不匹配的错误。
• get_single_element()方案：该方法是tf.data.Dataset的实例方法，仅能用于实际数据集对象，无法作用于符号化的KerasTensor，因此报错。

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