Застрял с реализацией самоконтроля

Question

Реализация самоконтроля в тензорном потоке Keras с битовой модификацией (например, остаточное (добавить соединение)).

У меня есть следующая форма ввода:

мой ввод: KerasTensor(type_spec=TensorSpec(shape=(None, 8, 6, 64), dtype=tf.float32, name=None), name='multiply/mul:0', description="created by layer 'multiply'")

Моя цель - обработать TensorSpec(shape=(None, 8, 6, 64)( 8 меток времени одну за другой ( 6 * 64 )) с помощью внимания к себе и получить карту функций внимания к себе для каждой метки времени, а затем снова объединить ее в форму выходного тензора (None, 8, 6, 64) .

Мой код для реализации самостоятельного внимания:

      def conv1d(x, channels, ks=1, strides=1, padding='same'):
    conv = tf.keras.layers.Conv1D(channels, ks, strides, padding, activation='relu', use_bias=False,
                                  kernel_initializer='HeNormal')(x)
    return conv

# Self attention
def my_self_attention(x, channels):  
    f = conv1d(x, channels)  # KerasTensor(type_spec=TensorSpec(shape=(None, 6, 64), dtype=tf.float32, name=None), name='conv1d/Relu:0', description="created by layer 'conv1d'")
    g = conv1d(x, channels) # KerasTensor(type_spec=TensorSpec(shape=(None, 6, 64), dtype=tf.float32, name=None), name='conv1d_1/Relu:0', description="created by layer 'conv1d_1'")
    h = conv1d(x, channels) # KerasTensor(type_spec=TensorSpec(shape=(None, 6, 64), dtype=tf.float32, name=None), name='conv1d_2/Relu:0', description="created by layer 'conv1d_2'")
    attention_weights = tf.keras.activations.softmax(
        tf.matmul(g, Permute((2, 1))(f)))  # KerasTensor(type_spec=TensorSpec(shape=(None, 6, 6), dtype=tf.float32, name=None), name='tf.math.truediv/truediv:0', description="created by layer 'tf.math.truediv'")
    sensor_att_fm = tf.matmul(attention_weights, h) # KerasTensor(type_spec=TensorSpec(shape=(None, 6, 64), dtype=tf.float32, name=None), name='tf.linalg.matmul_1/MatMul:0', description="created by layer 'tf.linalg.matmul_1'") 
    gamma = tf.compat.v1.get_variable("gamma", [1], initializer=tf.constant_initializer(0.0)) # <tf.Variable 'gamma:0' shape=(1,) dtype=float32, numpy=array([0.], dtype=float32)>  
    # to refined feature map adding feature map with input 
    o = gamma * sensor_att_fm + x # KerasTensor(type_spec=TensorSpec(shape=(None, 6, 64), dtype=tf.float32, name=None), name='tf.__operators__.add/AddV2:0', description="created by layer 'tf.__operators__.add'")
    return o

# Calling function 
sa = [my_self_attention(my_input[:, t, :, :], channels) for t in range(my_input.shape[1])]

Где я застрял

Я пытаюсь объединить их одну за другой на карте функций внимания и получить выходной тензор (None, 8, 6, 64), но не смог добиться успеха. Не могли бы вы помочь мне решить эту проблему. Во-вторых, убедитесь, что функции (my_self_attention и conv1d), которые я реализовал, в порядке.

Спасибо