Tensorflow 层在 model 外部工作但不在内部

Question

我有一个自定义的 tensorflow 层，通过生成 output 可以正常工作，但是当与 Keras 功能 model 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)

这会生成以下 output：

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

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

但是当我在 Keras 功能 API 中使用它时，它不起作用。 这是代码：

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. See https://www.tensorflow.org/api_docs/python/tf/dtypes for supported TF dtypes.

有谁知道为什么会发生这种情况以及如何解决？

Answer 1

我认为您应该尝试在自定义层中使用tf.shape ，因为它会给您张量的动态形状：

batch_size = tf.shape(x)[0]