TensorFlow自定义编解码器模型输出维度缺失问题排查

我构建了一个包含自定义encoder层和decoder层的编解码器模型，当前设置批量大小为128、序列长度为70、特征数为73，期望输出形状为128×70×73，但decoder层返回的是缺失特征维度的128×70二维输出。以下是编码器、解码器类定义、模型构建及训练代码，请问问题出在哪里？

原代码 #

编码器和解码器类定义 #

class encoder_layer(tf.keras.layers.Layer):
  def __init__(self):
    super(encoder_layer,self).__init__()

    # encoder block
    self.enc_input = Input(shape=(seq_length),name='enc_input')
    self.embed_layer = Embedding(input_dim=len(char2idx),output_dim=EMBEDDING_DIM,
                            weights=[embed_matrix],trainable=False,mask_zero=True)
    self.enc_lstm = Bidirectional(LSTM(lstm_size,return_state=True,return_sequences=True,name='enc_lstm'))

  def call(self,input_tensor,training=False):
    
    embed_data = self.embed_layer(input_tensor)
    enc_out,enc_state_h_f,enc_state_h_r,enc_state_c_f,enc_state_c_r = self.enc_lstm(embed_data)
    enc_state_h = tf.concat([enc_state_h_f,enc_state_h_r],axis=1)
    enc_state_c = tf.concat([enc_state_c_f,enc_state_c_r],axis=1)
    enc_state_h = tf.expand_dims(enc_state_h,axis = 1)
    enc_state_c = tf.expand_dims(enc_state_c,axis = 1)
    # self.states = [self.enc_state_h,self.enc_state_c]
    return tf.concat([enc_out,enc_state_h,enc_state_c],axis=1)


class decoder_layer(tf.keras.layers.Layer):
  def __init__(self):
    super(decoder_layer,self).__init__()
    # self.dec_input_layer = Input(shape=(1,seq_length+latentSpaceDimension))
    self.dec_lstm = LSTM(lstm_size*2,return_state=True,name='dec_lstm')

    self.attention = cust_attention(latentSpaceDimension)
    self.dec_dropout = Dropout(0.5)
    self.decoder_dense = Dense(num_features, activation='softmax',  name='decoder_dense')
    self.input_slice = Lambda(lambda x: tf.split(x, [-1,1,1],axis=1))
    self.final_dec_out = Lambda(lambda x: K.concatenate(x, axis=1))

  def compute_output_shape(self,input_shape):
    return(self.cur_batch_size,seq_length,len(char2idx))

  def call(self,input_tensor,training=False):
    # print('decoder input tensor shape:',input_tensor.shape)
    self.cur_batch_size = input_tensor.shape[0]
    # dec_input = self.dec_input_layer
    enc_out,enc_state_h,enc_state_c = self.input_slice(input_tensor)
    # print('enc_out shape - {},enc_state_h shape - {},enc_state_c shape - {}'\\
    #       .format(enc_out.shape,enc_state_h.shape,enc_state_c.shape))
    if self.cur_batch_size == None:
      return tf.zeros((1,seq_length))
    enc_state_h = tf.squeeze(enc_state_h,axis=1)
    enc_state_c = tf.squeeze(enc_state_c,axis=1)
    dec_in_data = np.zeros((self.cur_batch_size,1,num_features)).astype(np.float32)
    states = [enc_state_h,enc_state_c]
    all_outputs = []
    # print('enc_out shape - {},enc_state_h shape - {},enc_state_c shape - {}'\\
    #       .format(enc_out.shape,enc_state_h.shape,enc_state_c.shape))

    dec_out = enc_state_h

    for _ in range(seq_length):
      # print('******iter {}******'.format(_))

      context_vector, attention_weights = self.attention(dec_out,enc_out)
      context_vector = tf.expand_dims(context_vector,1)
      # if cur_batch_size != None:
      # print('pre concat context vector - {}, dec in data  - {}'\\
      #       .format(context_vector.shape,dec_in_data.shape))
      dec_in_data = tf.concat([context_vector,dec_in_data],axis=-1)
      # print('post concat dec in data  - {}'.format(dec_in_data.shape))
    
      dec_out,dec_state_h,dec_state_c = self.dec_lstm(dec_in_data,initial_state=states)

      states = [dec_state_h,dec_state_c]
      dense_output = self.decoder_dense(self.dec_dropout(dec_out))
      dense_output = tf.expand_dims(dense_output,1)
      dec_in_data = dense_output
      all_outputs.append(dense_output)
      print('output shape:',dense_output.shape)
    # print('output shape:',out_val.shape)
    # if cur_batch_size != None:
    print('len all outputs:',len(all_outputs))
    print('output shape:',self.final_dec_out(all_outputs).shape)
    return self.final_dec_out(all_outputs)

模型代码 #

def make_model():
  enc_layer = encoder_layer()
  dec_layer = decoder_layer()

  input_layer = Input(shape=(MAX_SENT_LEN+10))
  x = enc_layer(input_layer)
  output_layer = dec_layer(x)

  model = models.Model(inputs=input_layer,outputs=output_layer)
  return model

训练代码 #

x_model = make_model()
x_model.compile(optimizer=tf.keras.optimizers.RMSprop(),loss=tf.keras.losses.CategoricalCrossentropy(),metrics='accuracy')

history = x_model.fit(x=x_train,y = y_train,validation_data=(x_val,y_val),epochs=120,callbacks=callbacks_list,batch_size=batch_size)

问题核心出在解码器的模型构建阶段输出形状错误，以及几个细节上的张量操作不规范，具体修正点如下：

1. 解码器批量大小为None时的返回值错误 #

模型构建阶段（未传入实际数据时），cur_batch_size为None，此时原代码返回tf.zeros((1,seq_length))（二维张量），导致Keras错误推断输出形状为(None, seq_length)，而非期望的三维。需改为返回三维占位张量。

2. compute_output_shape方法定义错误 #

原方法中用运行时的cur_batch_size和错误的特征维度len(char2idx)定义输出形状，需改为基于输入形状推导的固定格式。

3. 其他细节修正 #

• 编码器中多余的Input定义移除，自定义Layer无需内部声明Input；
• 解码器中用tf.zeros替代np.zeros初始化张量，避免numpy与TensorFlow张量的兼容性问题；
• 解码器的Dropout层需明确传入training参数，保证训练/推理模式的行为一致；
• 模型构建函数改为传入参数，避免依赖全局变量，提升代码可维护性。

修改后的代码 #

编码器类 #

class encoder_layer(tf.keras.layers.Layer):
  def __init__(self, seq_length, lstm_size, char2idx, embed_matrix, EMBEDDING_DIM):
    super(encoder_layer,self).__init__()
    self.seq_length = seq_length
    self.lstm_size = lstm_size
    self.embed_layer = Embedding(input_dim=len(char2idx),output_dim=EMBEDDING_DIM,
                            weights=[embed_matrix],trainable=False,mask_zero=True)
    self.enc_lstm = Bidirectional(LSTM(lstm_size,return_state=True,return_sequences=True,name='enc_lstm'))

  def call(self,input_tensor,training=False):
    embed_data = self.embed_layer(input_tensor)
    enc_out,enc_state_h_f,enc_state_h_r,enc_state_c_f,enc_state_c_r = self.enc_lstm(embed_data)
    enc_state_h = tf.concat([enc_state_h_f,enc_state_h_r],axis=1)
    enc_state_c = tf.concat([enc_state_c_f,enc_state_c_r],axis=1)
    enc_state_h = tf.expand_dims(enc_state_h,axis = 1)
    enc_state_c = tf.expand_dims(enc_state_c,axis = 1)
    return tf.concat([enc_out, enc_state_h, enc_state_c], axis=1)

  def compute_output_shape(self, input_shape):
    return (input_shape[0], self.seq_length + 2, self.lstm_size * 2)

解码器类 #

class decoder_layer(tf.keras.layers.Layer):
  def __init__(self, seq_length, lstm_size, num_features, latentSpaceDimension):
    super(decoder_layer,self).__init__()
    self.seq_length = seq_length
    self.lstm_size = lstm_size
    self.num_features = num_features
    self.dec_lstm = LSTM(lstm_size*2,return_state=True,name='dec_lstm')

    self.attention = cust_attention(latentSpaceDimension)
    self.dec_dropout = Dropout(0.5)
    self.decoder_dense = Dense(num_features, activation='softmax',  name='decoder_dense')
    self.input_slice = Lambda(lambda x: tf.split(x, [-1,1,1],axis=1))
    self.final_dec_out = Lambda(lambda x: tf.concat(x, axis=1))

  def compute_output_shape(self, input_shape):
    return (input_shape[0], self.seq_length, self.num_features)

  def call(self,input_tensor,training=False):
    cur_batch_size = input_tensor.shape[0]
    enc_out,enc_state_h,enc_state_c = self.input_slice(input_tensor)
    
    if cur_batch_size is None:
      return tf.zeros((1, self.seq_length, self.num_features), dtype=tf.float32)
    
    enc_state_h = tf.squeeze(enc_state_h,axis=1)
    enc_state_c = tf.squeeze(enc_state_c,axis=1)
    dec_in_data = tf.zeros((cur_batch_size,1,self.num_features), dtype=tf.float32)
    states = [enc_state_h,enc_state_c]
    all_outputs = []

    dec_out = enc_state_h

    for _ in range(self.seq_length):
      context_vector, attention_weights = self.attention(dec_out,enc_out)
      context_vector = tf.expand_dims(context_vector,1)
      dec_in_data = tf.concat([context_vector,dec_in_data],axis=-1)
    
      dec_out,dec_state_h,dec_state_c = self.dec_lstm(dec_in_data,initial_state=states)

      states = [dec_state_h,dec_state_c]
      dense_output = self.decoder_dense(self.dec_dropout(dec_out, training=training))
      dense_output = tf.expand_dims(dense_output,1)
      dec_in_data = dense_output
      all_outputs.append(dense_output)
    
    return self.final_dec_out(all_outputs)

模型构建函数 #

def make_model(seq_length, lstm_size, char2idx, embed_matrix, EMBEDDING_DIM, num_features, latentSpaceDimension):
  enc_layer = encoder_layer(seq_length, lstm_size, char2idx, embed_matrix, EMBEDDING_DIM)
  dec_layer = decoder_layer(seq_length, lstm_size, num_features, latentSpaceDimension)

  input_layer = Input(shape=(seq_length))
  x = enc_layer(input_layer)
  output_layer = dec_layer(x)

  model = tf.keras.models.Model(inputs=input_layer,outputs=output_layer)
  return model

训练代码调用示例 #

# 假设已定义以下变量：
# seq_length=70, lstm_size=64, char2idx=xxx, embed_matrix=xxx, EMBEDDING_DIM=128, num_features=73, latentSpaceDimension=64

x_model = make_model(seq_length, lstm_size, char2idx, embed_matrix, EMBEDDING_DIM, num_features, latentSpaceDimension)
x_model.compile(optimizer=tf.keras.optimizers.RMSprop(),loss=tf.keras.losses.CategoricalCrossentropy(),metrics='accuracy')

history = x_model.fit(x=x_train,y = y_train,validation_data=(x_val,y_val),epochs=120,callbacks=callbacks_list,batch_size=128)

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