Transformer聊天机器人训练报错：Softmax层维度不匹配求助

尝试用Transformer搭建聊天机器人，训练过程中触发维度不匹配的ValueError，已尝试统一编码器与解码器的序列长度，但问题仍未解决。以下是完整代码、报错信息及数据维度详情：

完整代码 #

import pandas as pd
import re
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

INP_VOCAB_SIZE = 23212 + 2 # start, end tokens
TAR_VOCAB_SIZE = 23485 + 2 # start, end tokens

data = pd.read_csv("data.csv")

inputs = data["text"]
targets = data["answer"]

for idx in range(len(inputs)):
  inputs[idx] = "<start>" + inputs[idx] + "<end>"
  targets[idx] = "<start>" + targets[idx] + "<end>"

def get_vectorizer(inp, tar) -> tuple:
    input_vectorizer = layers.TextVectorization(
        INP_VOCAB_SIZE,
        output_sequence_length = 250,
        standardize=None,
    )

    target_vectorizer = layers.TextVectorization(
        TAR_VOCAB_SIZE,
        output_sequence_length = 451,
        standardize=None,
    )
    input_vectorizer.adapt(list(inp))
    target_vectorizer.adapt(list(tar))
    return input_vectorizer, target_vectorizer

input_vec_layer, tar_vec_layer = get_vectorizer(inputs, targets)

X = input_vec_layer(inputs)
outputs = tar_vec_layer(targets)
decoder_inputs = outputs[:, :-1]
y = outputs[:, 1:]

TRAIN_SIZE = 0.8

index = int(TRAIN_SIZE * len(X))

X_train = X[:index]
X_test = X[index:]
y_train = y[:index]
y_test = y[index:]
decoder_inputs_train = decoder_inputs[:index]
decoder_inputs_test = decoder_inputs[index:]

class TransformerEncoder(layers.Layer):
    def __init__(self, embed_dim, dense_dim, num_heads):
        super(TransformerEncoder, self).__init__()
        self.embed_dim = embed_dim
        self.dense_dim = dense_dim
        self.num_heads = num_heads
        self.attention = layers.MultiHeadAttention(
            num_heads=num_heads, key_dim=embed_dim
        )
        self.dense_proj = keras.Sequential([
            layers.Dense(dense_dim, activation="relu"),
            layers.Dense(embed_dim)
        ])
        self.layernorm_1 = layers.LayerNormalization()
        self.layernorm_2 = layers.LayerNormalization()
        self.supports_masking = True

    def call(self, inputs, mask=None):
        if mask is not None:
            padding_mask = tf.cast(mask[:, tf.newaxis, tf.newaxis, :], dtype="int32")

        attention_output = self.attention(
            query=inputs, value=inputs, key=inputs, attention_mask=padding_mask
        )
        proj_input = self.layernorm_1(inputs + attention_output)
        proj_output = self.dense_proj(proj_input)
        return self.layernorm_2(proj_input + proj_output)
      
class PositionalEmbedding(layers.Layer):
    def __init__(self, sequence_length, vocab_size, embed_dim):
        super(PositionalEmbedding, self).__init__()
        self.token_embeddings = layers.Embedding(
            input_dim=vocab_size, output_dim=embed_dim
        )
        self.position_embeddings = layers.Embedding(
            input_dim=sequence_length, output_dim=embed_dim
        )
        self.sequence_length = sequence_length
        self.vocab_size = vocab_size
        self.embed_dim = embed_dim

    def call(self, inputs):
        length = tf.shape(inputs)[-1]
        positions = tf.range(length)
        embedded_tokens = self.token_embeddings(inputs)
        embedded_positions = self.position_embeddings(positions)
        return embedded_tokens + embedded_positions

    def compute_mask(self, inputs, mask=None):
        return tf.math.not_equal(inputs, 0)

class TransformerDecoder(layers.Layer):
    def __init__(self, embed_dim, latent_dim, num_heads):
        super(TransformerDecoder, self).__init__()
        self.embed_dim = embed_dim
        self.latent_dim = latent_dim
        self.num_heads = num_heads
        self.attention_1 = layers.MultiHeadAttention(
            num_heads=num_heads, key_dim=embed_dim
        )
        self.attention_2 = layers.MultiHeadAttention(
            num_heads=num_heads, key_dim=embed_dim
        )
        self.dense_proj = keras.Sequential([
            layers.Dense(latent_dim, activation="relu"),
            layers.Dense(embed_dim)
        ])
        self.layernorm_1 = layers.LayerNormalization()
        self.layernorm_2 = layers.LayerNormalization()
        self.layernorm_3 = layers.LayerNormalization()
        self.supports_masking = True

    def call(self, inputs, encoder_outputs, mask=None):
        causal_mask = self.get_causal_attention_mask(inputs)
        if mask is not None:
            padding_mask = tf.cast(mask[:, tf.newaxis, :], dtype="int32")
            padding_mask = tf.minimum(padding_mask, causal_mask)

        attention_output_1 = self.attention_1(
            query=inputs, value=inputs, key=inputs, attention_mask=causal_mask
        )
        out_1 = self.layernorm_1(inputs + attention_output_1)

        attention_output_2 = self.attention_2(
            query=out_1, value=encoder_outputs,key=encoder_outputs,
            attention_mask=padding_mask
        )
        out_2 = self.layernorm_2(out_1 + attention_output_2)

        proj_output = self.dense_proj(out_2)
        return self.layernorm_3(out_2 + proj_output)

    def get_causal_attention_mask(self, inputs):
        input_shape = tf.shape(inputs)
        batch_size, sequence_length = input_shape[0], input_shape[1]
        i = tf.range(sequence_length)[:, tf.newaxis]
        j = tf.range(sequence_length)
        print(i, j)
        mask = tf.cast(i >= j, dtype="int32")
        mask = tf.reshape(mask, (1, input_shape[1], input_shape[1]))
        mult = tf.concat(
            [tf.expand_dims(batch_size, -1), tf.constant([1, 1], dtype=tf.int32)],
            axis=0,
        )
        return tf.tile(mask, mult)

def create_model():
    # encoder
    encoder_inputs = keras.Input(shape=(None,), dtype="int64", name="encoder_inputs")
    encoder_augmented_inputs = PositionalEmbedding(sequence_length = 250, vocab_size = INP_VOCAB_SIZE, embed_dim = 256)(encoder_inputs)
    encoder_outputs = TransformerEncoder(embed_dim = 256, dense_dim = 1024, num_heads = 8)(encoder_augmented_inputs)

    # decoder
    decoder_inputs = keras.Input(shape=(None,), dtype="int64", name="decoder_inputs")
    encoded_seq_inputs = keras.Input(shape=(None, 256), name="decoder_state_inputs") # embed dim = 256
    decoder_augmented_inputs = PositionalEmbedding(sequence_length = 450, vocab_size = TAR_VOCAB_SIZE, embed_dim = 256)(decoder_inputs)
    decoder_outputs = TransformerDecoder(embed_dim = 256, latent_dim = 1024, num_heads = 8)(decoder_augmented_inputs, encoded_seq_inputs)
    decoder_outputs = layers.Dropout(0.5)(decoder_outputs)
    decoder_outputs = layers.Dense(TAR_VOCAB_SIZE, activation="softmax")(decoder_outputs)

    decoder = keras.Model([decoder_inputs, encoded_seq_inputs], decoder_outputs, name='outputs')
    decoder_outputs = decoder([decoder_inputs, encoder_outputs])

    model = keras.Model(
        [encoder_inputs, decoder_inputs], decoder_outputs, name="transformer"
    )

    model.compile(
        "adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"]
    )

    return model

model = create_model()

history = model.fit(
  [X_train, decoder_inputs_train],
  y_train,
  validation_data=([X_test, decoder_inputs_test], y_test),
  epochs= 200,
  batch_size = 64
)

报错信息 #

ValueError                                Traceback (most recent call last)
<ipython-input-36-b682566c047c> in <cell line: 1>()
----> 1 history = model.fit(
      2   [X_train, decoder_inputs_train],
      3   y_train,
      4   validation_data=([X_test, decoder_inputs_test], y_test),
      5   epochs= 200,

/tmp/__autograph_generated_file4l44nhmd.py in tf__call(self, inputs, mask)
     25                 padding_mask = ag__.Undefined('padding_mask')
     26                 ag__.if_stmt(ag__.ld(mask) is not None, if_body, else_body, get_state, set_state, ('padding_mask',), 1)
---> 27                 attention_output = ag__.converted_call(ag__.ld(self).attention, (), dict(query=ag__.ld(inputs), value=ag__.ld(inputs), key=ag__.ld(inputs), attention_mask=ag__.ld(padding_mask)), fscope)
     28                 proj_input = ag__.converted_call(ag__.ld(self).layernorm_1, (ag__.ld(inputs) + ag__.ld(attention_output),), None, fscope)
     29                 proj_output = ag__.converted_call(ag__.ld(self).dense_proj, (ag__.ld(proj_input),), None, fscope)

ValueError: in user code:

    File "/usr/local/lib/python3.10/dist-packages/keras/engine/training.py", line 1284, in train_function  *
        return step_function(self, iterator)
    File "/usr/local/lib/python3.10/dist-packages/keras/engine/training.py", line 1268, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/usr/local/lib/python3.10/dist-packages/keras/engine/training.py", line 1249, in run_step  **
        outputs = model.train_step(data)
    File "/usr/local/lib/python3.10/dist-packages/keras/engine/training.py", line 1050, in train_step
        y_pred = self(x, training=True)
    File "/usr/local/lib/python3.10/dist-packages/keras/utils/traceback_utils.py", line 70, in error_handler
        raise e.with_traceback(filtered_tb) from None
    File "/tmp/__autograph_generated_file4l44nhmd.py", line 27, in tf__call
        attention_output = ag__.converted_call(ag__.ld(self).attention, (), dict(query=ag__.ld(inputs), value=ag__.ld(inputs), key=ag__.ld(inputs), attention_mask=ag__.ld(padding_mask)), fscope)

    ValueError: Exception encountered when calling layer 'transformer_encoder' (type TransformerEncoder).
    
    in user code:
    
        File "<ipython-input-25-d6821aecfdb2>", line 22, in call  *
            attention_output = self.attention(
        File "/usr/local/lib/python3.10/dist-packages/keras/utils/traceback_utils.py", line 70, in error_handler  **
            raise e.with_traceback(filtered_tb) from None
    
        ValueError: Exception encountered when calling layer 'softmax' (type Softmax).
        
        Dimensions must be equal, but are 8 and 64 for '{{node transformer/transformer_encoder/multi_head_attention/softmax/add}} = AddV2[T=DT_FLOAT](transformer/transformer_encoder/multi_head_attention/einsum/Einsum, transformer/transformer_encoder/multi_head_attention/softmax/mul)' with input shapes: [64,8,275,275], [64,64,275,275].
        
        Call arguments received by layer 'softmax' (type Softmax):
          • inputs=tf.Tensor(shape=(64, 8, 275, 275), dtype=float32)
          • mask=tf.Tensor(shape=(64, 64, 275, 275), dtype=bool)
    
    
    Call arguments received by layer 'transformer_encoder' (type TransformerEncoder):
      • inputs=tf.Tensor(shape=(64, 275, 256), dtype=float32)
      • mask=tf.Tensor(shape=(64, 275), dtype=bool)

数据维度 #

执行以下代码：

print(X_train.shape)
print(X_test.shape)
print(y_train.shape)
print(y_test.shape)
print(decoder_inputs_train.shape)
print(decoder_inputs_test.shape)

输出：

(40000, 275)
(10000, 275)
(40000, 450)
(10000, 450)
(40000, 450)
(10000, 450)

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