如何在for循环中使用Keras Layer的动态None维度解决形状转换报错

错误根因 #

• 你在build方法中错误将batch_input_shape（静态TensorShape对象）当做张量取维度，且将动态batch大小存入类固定属性，与输入的动态batch维度（None）冲突。
• 使用Python原生for循环遍历batch维度：Python循环会在计算图构建阶段直接展开，此时batch大小未确定，无法执行循环逻辑，直接触发形状转换报错。

修正方案 #

1. 移除类属性中存储固定batch大小的逻辑，运行时动态从输入张量中获取batch维度。
2. 用TensorFlow内置的tf.map_fn完成batch维度的遍历，该算子支持动态形状，会在运行时根据实际batch大小执行循环，无需在图构建阶段确定循环次数。内层遍历特征的循环次数是固定值（n_input=16），图构建阶段可正常展开，无需修改。
3. 修正compute_output_shape逻辑，直接返回与输入对齐的动态batch维度即可，无需绑定固定batch大小。

修正后代码 #

import tensorflow as tf
from tensorflow.keras.layers import Layer

class WeightedLayer(Layer):
    def __init__(self, n_input, n_memb, **kwargs):
        super(WeightedLayer, self).__init__(**kwargs)
        self.n = n_input   # 16 features
        self.m = n_memb    # 3 
        self.output_dim = self.m ** self.n # 提前计算输出特征维度，固定值
    
    def build(self, batch_input_shape):
        # 该层无自定义可训练参数，直接调用父类build即可
        super(WeightedLayer, self).build(batch_input_shape)
    
    def call(self, input_):
        # 封装单样本处理逻辑，input_per_sample形状为(16,3)
        def process_single_sample(input_per_sample):
            cp = input_per_sample[0, :]
            xd_shape = [self.m]
            c_shape = [1]
            for d in range(1, self.n):
                c_shape.insert(0, self.m)
                xd_shape.insert(0, 1)
                xd = tf.reshape(input_per_sample[d, :], xd_shape)
                c = tf.reshape(cp, c_shape)
                cp = tf.matmul(c, xd)
            return tf.reshape(cp, (self.output_dim,))
        
        # 用tf.map_fn批量处理所有样本，自动适配动态batch大小
        return tf.map_fn(process_single_sample, input_, dtype=tf.float32)

    def compute_output_shape(self, batch_input_shape):
        # 输出batch维度与输入一致，特征维度为提前计算的固定值
        return tf.TensorShape([batch_input_shape[0], self.output_dim])

适配原有训练逻辑 #

你原有模型构建、编译、训练的代码无需修改，直接使用修正后的WeightedLayer即可正常运行。

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