умножение тензоров, когда одно из измерений не является постоянным

Question

У меня есть модель BERT, и я пытаюсь отфильтровать некоторые слова (временные шаги) из вывода. Это выходная форма слоя BERT:

Tensor("bert_layer/bert_layer_module_apply_tokens/bert/encoder/Reshape_13:0", shape=(?, ?, 768), dtype=float32)

, где первый ? - для batch_size, а второй - для количества временных шагов.

Я хочу применить умножение матриц между выводом BERT и вектором нулей и единиц в позициях слов, которые я хочу сохранить (например, 0,0,0,0,1,0,1,0,0 et c.). Таким образом, размерность этого вектора - это количество временных шагов.

Tensor("ExpandDims:0", shape=(?, 200), dtype=float32)

Чтобы умножить, я делаю следующее (см. filtering функция):

bert_max_seq_length = 200
in_id = tf.keras.layers.Input(shape=(bert_max_seq_length,), name="input_ids")
in_mask = tf.keras.layers.Input(shape=(bert_max_seq_length,), name="input_masks")
in_segment = tf.keras.layers.Input(shape=(bert_max_seq_length,), name="segment_ids")

# The filtering vector
e1_mask = tf.keras.layers.Input(shape=(bert_max_seq_length,), name="e1_masks")

bert_inputs = [in_id, in_mask, in_segment]

net = BertLayer(n_fine_tune_layers=0, bert_model='bert-base-uncased')(bert_inputs)
def filtering(sequence_output, e_mask):
    extended_e_mask = tf.keras.backend.expand_dims(e_mask, 1) # `(A)`
    # extended_e_mask shape becomes (?, 1, 200)
    dotted = tf.keras.layers.dot([e_mask, sequence_output], axes=[1, 1]) # `(B)`
    return dotted
e1_h = filtering(net, e1_mask)

pred = tf.keras.layers.Dense(2, activation="softmax")(e1_h)
model = tf.keras.models.Model(inputs=[*bert_inputs, e1_mask], outputs=pred)
optimizer = tf.keras.optimizers.Adam(lr=0.001)
model.compile(loss="categorical_crossentropy", optimizer=optimizer, metrics=["accuracy"])
model.summary()

В точке A Я добавил еще одно измерение к фильтрующему вектору, чтобы он был в форме (1x200). А в точке B я хочу применить умножение матриц между матрицами (?,1,200) и (?,?, 765). Я передаю [1,1] параметру axes, чтобы указать размеры, в которых должен выполняться процесс умножения.

Проблема в том, что вывод BERT не определен во втором измерении, и поэтому я получаю эту ошибку :

ValueError: Dimension incompatibility 200 != None. Layer shapes: (None, 200), (None, None, 768)

Как я могу решить эту проблему?