实现self attention简单模型

以下是一个简单的self attention模型的实现：

import tensorflow as tf

class SelfAttention(tf.keras.layers.Layer):
    def __init__(self, units):
        super(SelfAttention, self).__init__()
        self.W1 = tf.keras.layers.Dense(units)
        self.W2 = tf.keras.layers.Dense(units)
        self.V = tf.keras.layers.Dense(1)

    def call(self, inputs):
        query = self.W1(inputs)
        key = self.W2(inputs)
        score = tf.nn.tanh(tf.matmul(query, key, transpose_b=True))
        attention_weights = tf.nn.softmax(self.V(score), axis=1)
        context_vector = attention_weights * inputs
        context_vector = tf.reduce_sum(context_vector, axis=1)
        return context_vector

这个self attention模型由两个全连接层（W1和W2）和一个输出层（V）组成。输入经过W1和W2后得到query和key，然后计算score，再通过softmax层得到attention_weights，最后将attention_weights与输入相乘得到context_vector。最终返回context_vector作为输出。