自定义MultiHeadSelfAttention层在Keras函数式API中报错求修复

错误原因 #

在split_heads和merge_heads方法中，你使用静态形状获取方式x.shape[0]获取批量大小。当使用Keras函数式API的Input层时，输入的批量维度是None（表示批量大小未知），此时x.shape[0]返回None。在tf.reshape中同时存在None和-1会导致形状解析失败，触发类型转换错误。

而模拟数据测试时，输入是具体的numpy数组，x.shape[0]是明确数值（如1），因此不会出现问题。

修复方法 #

将获取批量大小的方式从静态形状改为动态形状，即用tf.shape(x)[0]替代x.shape[0]。tf.shape会在运行时获取张量的实际形状，无论构建模型时是否知道批量大小，都能正确解析。

修复后的关键代码 #

修改split_heads和merge_heads方法：

def split_heads(self, x):
    # 用tf.shape获取动态批量大小
    batch_size = tf.shape(x)[0]
    split_inputs = tf.reshape(x, (batch_size, -1, self.num_heads, self.d_head))
    return tf.transpose(split_inputs, perm=[0, 2, 1, 3])
  
def merge_heads(self, x):
    # 用tf.shape获取动态批量大小
    batch_size = tf.shape(x)[0]
    merged_inputs = tf.transpose(x, perm=[0, 2, 1, 3])
    return tf.reshape(merged_inputs, (batch_size, -1, self.d_model))

完整修复后的代码 #

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

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

  if mask is not None:
    scaled_scores = tf.where(mask==0, -np.inf, 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_size = tf.shape(x)[0]
    split_inputs = tf.reshape(x, (batch_size, -1, self.num_heads, self.d_head))
    return tf.transpose(split_inputs, perm=[0, 2, 1, 3])
  
  def merge_heads(self, x):
    batch_size = tf.shape(x)[0]
    merged_inputs = tf.transpose(x, perm=[0, 2, 1, 3])
    return tf.reshape(merged_inputs, (batch_size, -1, self.d_model))

  def call(self, q, k, v, mask):
    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)

# ----- Testing with simulated data ------- 
x = np.random.rand(1,2,3)
values_emb = MultiHeadSelfAttention(3, 3)(x,x,x, mask = None)
print(values_emb)

# ----- Keras Functional API Test -------
x_input = Input(shape=(2,3))
values_emb = MultiHeadSelfAttention(3, 3)(x_input,x_input,x_input, mask = None)
model = Model(x_input, values_emb)
model.summary()

验证结果 #

修复后，函数式API可正常构建模型并打印summary，模拟数据测试也能保持原有输出。

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