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]