如何在TF.Keras中提取单一数据源的时间序列滑动窗口并解决形状错误？

我有如下格式的数据集：

time | src | a | b | c | d | e | Label
----------------------------------------
0.   | 1   | # | # | # | # | # | #
1.   | 1   | # | # | # | # | # | #
2.   | 1   | # | # | # | # | # | #
3.   | 1   | # | # | # | # | # | #
4.   | 1   | # | # | # | # | # | #
....
0.   | 2   | # | # | # | # | # | #
1.   | 2   | # | # | # | # | # | #
2.   | 2   | # | # | # | # | # | #
3.   | 2   | # | # | # | # | # | #
4.   | 2   | # | # | # | # | # | #

我正在训练模型，目标是用[a,b,c,d,e]的窗口数据预测窗口末尾的Label值。其中输入X的形状为(window_size,5)，输出y为窗口末尾的Label值，且要求X的所有数据必须来自同一src（即一个窗口仅包含单一数据源的数据）。

此前我通过tf.keras.utils.Sequence构建X/y对来管理内存，现在想改用tf.keras.utils.timeseries_dataset_from_array，但该工具无法识别src字段，可能生成跨数据源的窗口。请问如何利用该工具提取仅包含单一数据源的滑动窗口（即每个数据源的所有重叠窗口）？

进展 #

1. 未区分数据源的成功实现 #

我成功使用timeseries_dataset_from_array生成了数据集，但没有区分数据源：

# ============= 数据准备 ===========

import tensorflow as tf
import numpy as np

# 构造模拟数据
X = np.random.random((100,5))
y = np.random.random((100))
src = np.array([0]*50 + [1]*50)
window_size = 5

# 生成不区分数据源的时间序列数据集
Xy_ds = tf.keras.utils.timeseries_dataset_from_array(
    X, y, batch_size = 2, sequence_length=window_size, 
    sequence_stride=window_size, shuffle=True
)

# ============= 模型训练 ===========
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, InputLayer, LSTM, Flatten

# 构建并训练模型，验证数据集可正常工作
model = Sequential()
model.add(InputLayer(input_shape=[window_size,5]))
model.add(LSTM(3))
model.add(Flatten())
model.add(Dense(1, activation='relu'))
optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)

model.compile(loss='categorical_crossentropy', optimizer=optimizer, metrics=['accuracy'])
history = model.fit(Xy_ds,epochs=1)

2. 尝试过滤单一数据源但出现形状错误 #

我尝试了一种解决方案，但训练时出现形状不匹配的错误：

# ============= 数据准备 ===========

import tensorflow as tf
import numpy as np

# 构造模拟数据
X = np.random.random((100,5))
y = np.random.random((100))
src = np.expand_dims(np.array([0]*50 + [1]*50),1)
window_size = 5

# 将src信息附加到X中，用于后续过滤
X = np.append(src, X, 1)

# 生成不区分数据源的时间序列数据集
Xy_ds = tf.keras.utils.timeseries_dataset_from_array(
    X, y, sequence_length=window_size, 
    sequence_stride=1, shuffle=True
)

def single_source(x,y):
    source = x[:,0]
    return tf.reduce_all(source == source[0])
    
# 过滤后移除src信息
def drop_source(x,y):
    return x[:, 1:], y

Xy_ds = Xy_ds.filter(single_source).map(drop_source)

# ============= 模型训练 ===========
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, InputLayer, LSTM, Flatten

# 构建并训练模型
model = Sequential()
model.add(InputLayer(input_shape=[window_size,5]))
model.add(LSTM(3))
model.add(Flatten())
model.add(Dense(1, activation='relu'))
optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)

model.compile(loss='categorical_crossentropy', optimizer=optimizer, metrics=['accuracy'])
history = model.fit(Xy_ds,epochs=1)

错误信息：

ValueError: 输入0与层sequential_3不兼容：期望形状=(None, None, 5)，找到形状=(None, None, 6)

我尝试了相关的形状修复方案，但均未解决问题。

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