TensorFlow自定义层结合Functional API报错：无法转符号张量为NumPy数组

错误根源 #

报错核心是混用NumPy函数处理TensorFlow符号张量：在scaled_dot_product_attention函数中，np.sqrt(key_dim)和-np.inf属于NumPy操作，但key_dim是TensorFlow的符号张量（模型构建阶段未绑定具体数值），NumPy无法处理这类延迟计算的张量，导致转换失败。

修复步骤 #

将所有NumPy相关的张量操作替换为TensorFlow原生函数：

• 把np.sqrt(key_dim)改为tf.sqrt(key_dim)，用TensorFlow平方根函数处理张量
• 把-np.inf改为tf.constant(-np.inf, dtype=tf.float32)，用TensorFlow常量替代NumPy常量，保证张量类型一致

同时修正split_heads和merge_heads的维度处理逻辑，避免丢失序列长度维度，确保注意力计算的维度正确性。

修复后的完整代码 #

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

def scaled_dot_product_attention(query, key, value, mask=None):
  key_dim = tf.cast(tf.shape(key)[-1], tf.float32)
  # 替换np.sqrt为tf.sqrt
  scaled_scores = tf.matmul(query, key, transpose_b=True) / tf.sqrt(key_dim)

  if mask is not None:
    # 替换-np.inf为TensorFlow常量
    scaled_scores = tf.where(mask==0, tf.constant(-np.inf, dtype=tf.float32), scaled_scores)

  softmax = tf.keras.layers.Softmax()
  weights = softmax(scaled_scores) 
  return tf.matmul(weights, value), weights

class MultiHeadSelfAttention(tf.keras.layers.Layer):
  def __init__(self, d_model, num_heads):
    super(MultiHeadSelfAttention, self).__init__()
    self.d_model = d_model
    self.num_heads = num_heads

    self.d_head = self.d_model // self.num_heads

    self.wq = tf.keras.layers.Dense(self.d_model)
    self.wk = tf.keras.layers.Dense(self.d_model)
    self.wv = tf.keras.layers.Dense(self.d_model)

    # Linear layer to generate the final output.
    self.dense = tf.keras.layers.Dense(self.d_model)
  
  def split_heads(self, x):
    # 保留batch和seq_len维度，拆分多头
    split_inputs = tf.reshape(x, (-1, tf.shape(x)[1], self.num_heads, self.d_head))
    return tf.transpose(split_inputs, perm=[0,2,1,3])
  
  def merge_heads(self, x):
    merged_inputs = tf.transpose(x, perm=[0,2,1,3])
    return tf.reshape(merged_inputs, (-1, tf.shape(x)[2], self.d_model))

  def call(self, q, k, v, mask=None):
    qs = self.wq(q)
    ks = self.wk(k)
    vs = self.wv(v)

    qs = self.split_heads(qs)
    ks = self.split_heads(ks)
    vs = self.split_heads(vs)

    output, attn_weights = scaled_dot_product_attention(qs, ks, vs, mask)
    output = self.merge_heads(output)

    return self.dense(output) 

# -------- Model Building -----------
values = Input(shape=(3,12))
values_emb = MultiHeadSelfAttention(12, 3)(values, values, values, mask=None)
model = Model(values, values_emb)
model.summary()

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