Вход для изменения формы не соответствует запрошенной форме

Question

Я знаю, что другие уже отправляли похожие вопросы, но я не смог найти здесь подходящего решения.

Я написал собственный слой keras для усреднения результатов DistilBert на основе маски. То есть у меня естьdim=[batch_size, n_tokens_out, 768] входя в маску n_tokens_out на основе маски, которая dim=[batch_size, n_tokens_out]. Результат должен бытьdim=[batch_size, 768]. Вот код слоя:

class CustomPool(tf.keras.layers.Layer):
    def __init__(self, output_dim, **kwargs):
        self.output_dim = output_dim
        super(CustomPool, self).__init__(**kwargs)
    
    def call(self, x, mask):
        masked = tf.cast(tf.boolean_mask(x, mask = mask, axis = 0), tf.float32)
        mn = tf.reduce_mean(masked, axis = 1, keepdims=True)
        return tf.reshape(mn, (tf.shape(x)[0], self.output_dim))
    
    def compute_output_shape(self, input_shape):
        return (input_shape[0], self.output_dim)

Модель компилируется без ошибок, но как только начинается обучение, появляется такая ошибка:

InvalidArgumentError: 2 root error(s) found.
  (0) Invalid argument:  Input to reshape is a tensor with 967 values, but the requested shape has 12288
     [[node pooled_distilBert/CustomPooling/Reshape (defined at <ipython-input-245-a498c2817fb9>:13) ]]
     [[assert_greater_equal/Assert/AssertGuard/pivot_f/_3/_233]]
  (1) Invalid argument:  Input to reshape is a tensor with 967 values, but the requested shape has 12288
     [[node pooled_distilBert/CustomPooling/Reshape (defined at <ipython-input-245-a498c2817fb9>:13) ]]
0 successful operations.
0 derived errors ignored. [Op:__inference_train_function_211523]

Errors may have originated from an input operation.
Input Source operations connected to node pooled_distilBert/CustomPooling/Reshape:
 pooled_distilBert/CustomPooling/Mean (defined at <ipython-input-245-a498c2817fb9>:11)

Input Source operations connected to node pooled_distilBert/CustomPooling/Reshape:
 pooled_distilBert/CustomPooling/Mean (defined at <ipython-input-245-a498c2817fb9>:11)

Размеры, которые я получаю, меньше ожидаемых, что для меня странно.

Вот как выглядит модель (TFDistilBertModel взят из huggingface transformers библиотека):

dbert_layer = TFDistilBertModel.from_pretrained('distilbert-base-uncased')

in_id = tf.keras.layers.Input(shape=(seq_max_length,), dtype='int32', name="input_ids")
in_mask = tf.keras.layers.Input(shape=(seq_max_length,), dtype='int32', name="input_masks")
    
dbert_inputs = [in_id, in_mask]
dbert_output = dbert_layer(dbert_inputs)[0]
x = CustomPool(output_dim = dbert_output.shape[2], name='CustomPooling')(dbert_output, in_mask)
dense1 = tf.keras.layers.Dense(256, activation = 'relu', name='dense256')(x)
pred = tf.keras.layers.Dense(n_classes, activation='softmax', name='MODEL_OUT')(dense1)

model = tf.keras.models.Model(inputs = dbert_inputs, outputs = pred, name='pooled_distilBert')

Любая помощь здесь будет очень ценится, как я имел взгляд через существующие вопросы, большинство в конечном итоге решается путем указания формы ввода (не применимо в моем случае).

Answer 1

Использование tf.reshape перед слоем пула

Я знаю, что мой ответ запоздал, но я хочу поделиться своим решением проблемы. Дело в том, что вы пытаетесь изменить фиксированный размер вектора (тензора) во время обучения модели. Вектор изменит свой входной размер, а фиксированное изменение формы, например tf.reshape(updated_inputs, (shape = fixed_shape)) вызовет ваши проблемы, на самом деле моя проблема :)) Надеюсь, это поможет