自定义TensorFlow Layer单独可用但无法在Keras函数式模型中运行

我实现了一个自定义TensorFlow Layer，单独传入样本数据测试时可正常生成输出，但将其用于Keras函数式模型API时会抛出错误。

测试代码：

import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input

# ------ Custom Layer -----------
class CustomLayer(tf.keras.layers.Layer):
  def __init__(self):
    super(CustomLayer, self).__init__()

  def split_heads(self, x):
    
    batch_size = x.shape[0]
    split_inputs = tf.reshape(x, (batch_size, -1, 3, 1))
    
    return split_inputs
  
  def call(self, q):

    qs = self.split_heads(q)

    return qs

# ------ Testing Layer with sample data --------
x = np.random.rand(1,2,3)
values_emb = CustomLayer()(x)
print(values_emb)

测试输出：

tf.Tensor(
[[[[0.7148978 ]
   [0.3997009 ]
   [0.11451813]]

  [[0.69927174]
   [0.71329576]
   [0.6588452 ]]]], shape=(1, 2, 3, 1), dtype=float32)

使用Keras函数式模型时的问题代码：

x = Input(shape=(2,3))
values_emb = CustomLayer()(x)
model = Model(x, values_emb)
model.summary()

触发错误：

TypeError: Failed to convert elements of (None, -1, 3, 1) to Tensor. Consider casting elements to a supported type.

问题原因 #

问题出在静态形状与动态形状的差异：

• 单独测试时，输入是具体的numpy数组，x.shape[0]能拿到实际的批量大小（比如1），tf.reshape可以正常解析(1, -1, 3, 1)。
• 但在函数式模型中，Input的批量维度是None（表示批量大小可变），此时x.shape[0]返回的是None，而tf.reshape不支持(None, -1, 3, 1)这种包含两个不确定维度（None和-1）的写法——-1本身是让TensorFlow自动计算剩余维度，与None组合会导致形状无法确定，从而触发类型转换错误。

修复方法 #

用动态形状获取方式替代静态形状：使用tf.shape(x)[0]获取运行时的实际批量大小，或者直接用-1代替批量维度，让TensorFlow自动适配。

修复后的自定义Layer代码 #

class CustomLayer(tf.keras.layers.Layer):
  def __init__(self):
    super(CustomLayer, self).__init__()

  def split_heads(self, x):
    # 方式1：获取动态批量大小
    batch_size = tf.shape(x)[0]
    split_inputs = tf.reshape(x, (batch_size, -1, 3, 1))
    
    # 方式2：更简洁的写法，直接用-1适配批量维度
    # split_inputs = tf.reshape(x, (-1, tf.shape(x)[1], 3, 1))
    
    return split_inputs
  
  def call(self, q):
    qs = self.split_heads(q)
    return qs

验证函数式模型 #

运行以下代码即可正常生成模型并打印summary：

x = Input(shape=(2,3))
values_emb = CustomLayer()(x)
model = Model(x, values_emb)
model.summary()

输出示例：

Model: "model"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 input_1 (InputLayer)        [(None, 2, 3)]            0         
                                                                 
 custom_layer (CustomLayer)  (None, 2, 3, 1)           0         
                                                                 
=================================================================
Total params: 0
Trainable params: 0
Non-trainable params: 0
_________________________________________________________________

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