自定义Keras层输出形状为None致model.fit()报错的技术问询

你的问题核心在于Keras无法自动推断自定义层的输出形状，根源有两点：

• 错误重写output_shape属性，违背了Keras的形状推断逻辑
• call方法中对分母的维度操作（squeeze+expand_dims）干扰了TensorFlow的静态形状推断

修正步骤 #

1. 移除自定义的output_shape属性：Keras依赖compute_output_shape方法或张量静态形状完成推断，手动重写该属性会引发冲突。
2. 实现compute_output_shape方法：由于你的层不改变输入形状，直接返回输入形状即可。
3. 简化call方法中的分母计算：使用tf.reduce_sum的keepdims=True参数，避免不必要的维度增减，确保形状推断正常。

修正后的完整代码 #

import tensorflow as tf
from tensorflow import keras

class B_Softmax(keras.layers.Layer):
    def __init__(self, b_init_mean=10, b_init_var=0.001):
        super(B_Softmax, self).__init__()
        self.b_init = tf.random_normal_initializer(b_init_mean, b_init_var)
        
    def build(self, input_shape):
        self.b = tf.Variable(
            initial_value = self.b_init(shape=(1,), dtype='float32'),
            trainable=True
        )
        super().build(input_shape)  # 调用父类方法完成层初始化

    def call(self, inputs):
        nominators = tf.math.exp(self.b * inputs)
        # 使用keepdims=True保持维度，避免手动调整维度的操作
        denominator = tf.reduce_sum(nominators, axis=1, keepdims=True)
        s = tf.divide(nominators, denominator)
        return s

    def compute_output_shape(self, input_shape):
        # 直接返回输入形状，因为层不改变张量形状
        return input_shape

验证效果 #

1. 单独使用层：功能与之前一致：

A     = tf.constant([[1,2,3], [7,5,6]], dtype="float32")
layer = B_Softmax(1.0)
print(layer(A))

输出：

<tf.Tensor: shape=(2, 3), dtype=float32, numpy=
array([[0.08991686, 0.24461554, 0.6654676 ],
       [0.6654677 , 0.08991687, 0.24461551]], dtype=float32)>

2. 放入Sequential模型：model.summary()将正确显示输出形状：

input_dim = 5
num_classes = 10
model = keras.Sequential([
        keras.layers.Dense(32, activation='relu', input_shape=(input_dim,)),
        keras.layers.Dense(num_classes, activation="softmax"),
        B_Softmax(1.0)
])
model.summary()

输出：

Model: "sequential"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 dense (Dense)               (None, 32)                192       
                                                                 
 dense_1 (Dense)             (None, 10)                330       
                                                                 
 b__softmax (B_Softmax)      (None, 10)                1         
                                                                 
=================================================================
Total params: 523
Trainable params: 523
Non-trainable params: 0
_________________________________________________________________

3. 训练模型：model.fit()可正常运行，不再抛出ValueError: as_list() is not defined on an unknown TensorShape异常。

关键说明 #

• compute_output_shape是Keras自定义层中用于静态形状推断的标准方法，当输出形状可由输入形状直接推导时，必须实现该方法。
• keepdims=True让tf.reduce_sum保持输出张量的维度与输入一致，避免手动调整维度带来的形状推断问题。
• 禁止手动重写output_shape属性，这会干扰Keras内部的形状推断逻辑。

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