Keras自定义层获取输入形状报错：无法转换部分已知TensorShape

需求与初始实现 #

需要实现一个Keras自定义层MyLayer，要求前3次调用时返回与输入同形状的零张量，后续调用直接返回输入张量。初始代码如下：

class MyLayer(tf.keras.layers.Layer):

    def __init__(self, **kwargs):
        super(MyLayer, self).__init__(**kwargs)
        self.__iteration = 0
        self.__returning_zeros = None

    def build(self, input_shape):
        self.__returning_zeros = tf.zeros(shape=input_shape, dtype=tf.float32)

    def call(self, inputs):
        self.__iteration += 1

        if self.__iteration <= 3:
            return self.__returning_zeros
        else:
            return inputs

模型构建代码 #

将该层加入模型后，构建代码如下：

def build_model(input_shape, num_classes):
    input_layer = keras.Input(shape=input_shape, name='input')
    conv1 = layers.Conv2D(32, kernel_size=(3, 3), activation="relu", name='conv1')(input_layer)
    maxpool1 = layers.MaxPooling2D(pool_size=(2, 2), name='maxpool1')(conv1)
    conv2 = layers.Conv2D(64, kernel_size=(3, 3), activation="relu", name='conv2')(maxpool1)
    mylayer = MyLayer()(conv2)
    maxpool2 = layers.MaxPooling2D(pool_size=(2, 2), name='maxpool2')(mylayer)
    flatten = layers.Flatten(name='flatten')(maxpool2)
    dropout = layers.Dropout(0.5, name='dropout')(flatten)
    dense = layers.Dense(num_classes, activation="softmax", name='dense')(dropout)

    return keras.Model(inputs=(input_layer,), outputs=dense)

触发的错误 #

运行时出现以下错误：

File "customlayerkeras.py", line 25, in build
    self.__returning_zeros = tf.zeros(shape=input_shape, dtype=tf.float32)
ValueError: Cannot convert a partially known TensorShape (None, 13, 13, 64) to a Tensor.

错误原因 #

build方法中接收的input_shape包含动态维度（批量维度为None），因为Keras在模型构建阶段不确定实际输入的批量大小。tf.zeros无法直接将包含None的形状转换为张量，必须基于实际输入的具体形状生成零张量。

最优解决方案 #

无需在build中提前创建零张量，直接在call方法里基于输入生成零张量即可。修改后的call方法如下：

def call(self, inputs):
    self.__iteration += 1

    if self.__iteration <= 3:
        return inputs * 0
    else:
        return inputs

也可以用tf.zeros_like(inputs)替代inputs*0，效果完全一致：

def call(self, inputs):
    self.__iteration += 1

    if self.__iteration <= 3:
        return tf.zeros_like(inputs)
    else:
        return inputs

这两种方式都会根据每次输入的实际形状生成零张量，完美适配动态批量维度的场景。

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