TensorFlow 2.0 flat_map展平嵌套Dataset返回cardinality=-2及训练形状未知问题

参考TensorFlow官方文档的方法处理自有数据，按窗口规则切分生成嵌套Dataset后做展平操作，目前遇到的核心问题如下：

• 展平得到的flat_windows调用cardinality().numpy()获取Dataset基数时，返回值为-2，直接引发训练流程异常
• 尝试定位set_cardinality相关接口修正Dataset基数时，未找到可直接调用的可用实现；换用其他嵌套Dataset展平方案也均未跑通
• 上述问题传导到模型训练环节：采用子类化实现的自定义模型在图模式下训练时，内部的Linear层、Dense层会触发输入形状未知的报错；仅通过tf.config.run_functions_eagerly(True)开启eager模式时模型可正常运行，但训练速度过慢，最终目标是为图模式训练提供形状明确的输入数据。

窗口构建与展平的原始代码 #

window_size = 5
windows = range_ds.window(window_size, shift=1)
for sub_ds in windows.take(5):
    print(sub_ds)

flat_windows = windows.flat_map(lambda x: x)

class NeuralNetworkModel(tf.keras.Model): 
    def __init__(self):
        super(NeuralNetworkModel, self).__init__()
        self.encoder = Encoder()        
    
    def train_step(self, inputs):       
        X = inputs[0]
        Y = inputs[1] 
        
        with tf.GradientTape() as tape:
            enc_X = self.encoder(X)
            enc_Y = self.encoder(Y)    
            # 损失计算
            loss = tf.norm(enc_Y - enc_X, axis=[0, 1], ord='fro')
                
        # 计算梯度
        trainable_vars = self.encoder.trainable_variables
        gradients = tape.gradient(loss, trainable_vars)

        # 更新模型权重
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))

        # 更新自定义指标状态
        loss_tracker.update_state(loss)
        
        return {"loss": loss_tracker.result()}
        
    @property
    def metrics(self):
        # 注册自定义Metric对象，实现epoch开始/评估阶段自动重置状态
        return [loss_tracker]
    
    def test_step(self, inputs):       
        X = inputs[0]
        Y = inputs[1] 

        Psi_X = self.encoder(X)
        Psi_Y = self.encoder(Y)    

        # 损失计算
        loss = tf.norm(Psi_Y - Psi_X, axis=[0, 1], ord='fro')
        loss_tracker.update_state(loss)
        
        return {"loss": loss_tracker.result()}
        
class Encoder(tf.keras.Model):
    def __init__(self):
        super(Encoder, self).__init__(dtype='float64', name='Encoder')
        self.input_layer = DenseLayer(128)
        self.hidden_layer1 = DenseLayer(128)
        self.hidden_layer2 = DenseLayer(64)        
        self.hidden_layer3 = DenseLayer(64)
        self.output_layer = LinearLayer(64)
        
    def call(self, input_data, training):
        fx = self.input_layer(input_data)        
        fx = self.hidden_layer1(fx)
        fx = self.hidden_layer2(fx)
        fx = self.hidden_layer3(fx)
        return self.output_layer(fx)    

class LinearLayer(tf.keras.layers.Layer):
    def __init__(self, units):
        super(LinearLayer, self).__init__(dtype='float64')
        self.units = units

    def build(self, input_shape):
        input_dim = input_shape[-1]
        self.w = self.add_weight(
            shape=(input_dim, self.units), 
            initializer="random_normal", 
            trainable=True
        )
        self.b = self.add_weight(
            shape=(self.units,),    
            initializer=tf.zeros_initializer(),
            trainable=True
        )

    def call(self, inputs):
        return tf.matmul(inputs, self.w) + self.b

class DenseLayer(tf.keras.layers.Layer):
    def __init__(self, units):
        super(DenseLayer, self).__init__(dtype='float64')
        self.units = units
    
    def build(self, input_shape):
        input_dim = input_shape[-1]
        self.w = self.add_weight(
            shape=(input_dim, self.units), 
            initializer="random_normal", 
            trainable=True
        )
        self.b = self.add_weight(
            shape=(self.units,),    
            initializer=tf.zeros_initializer(),
            trainable=True
        )

    def call(self, inputs):
        x = tf.matmul(inputs, self.w) + self.b
        return tf.nn.elu(x)

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