MultHeadAttention.py

from mxnet.gluon import nn
from mxnet import nd
import copy
def softmax_attention(x):
    e=nd.exp(x-nd.max(x,axis=1,keepdims=True))
    s=nd.sum(e,axis=1,keepdims=True)
    return e/s
class Context_Attention(nn.Block):
    def __init__(self,**kwargs):
        super(Context_Attention, self).__init__(**kwargs)

        self.dense1_layer=nn.Dense(units=19,activation="tanh")

    def forward(self, maxlen_input,h,st_1):
        # print(self.collect_params())
        st=nd.reshape(st_1,shape=(st_1.shape[1],st_1.shape[2]))
        st1=nd.repeat(st,repeats= maxlen_input,axis=1)
        x=nd.concat(h,st1,dim=1)
        x=self.dense1_layer(x)
        alphas=softmax_attention(x)
        context=nd.dot(alphas,h)
        context=context.reshape(1,context.shape[0],context.shape[1])
        return context
# def scale_dot_product_attention(query,key,value,mask):
#
#     depth= key.shape[-1]
#     values=nd.dot(query,key,transpose_b=True)\
#            # /nd.sqrt(depth)
#     values=values/nd.array([depth],dtype="float32")
#     if  mask is not  None:
#         values+=mask*-1e9
#     attention_weights=nd.softmax(values,axis=-1)
#     output=nd.dot(attention_weights,value)
#     return output
# class MultHeadAttention(nn.Block):
#     def __init__(self,num_hiddens,num_heads,**kwargs):
#         super(MultHeadAttention, self).__init__(**kwargs)
#         self.query_dense=nn.Dense(num_hiddens,activation="sigmoid")
#         self.num_heads=num_heads
#         self.key_dense=nn.Dense(num_hiddens,activation="sigmoid")
#         self.value_dense=nn.Dense(num_hiddens,activation="sigmoid")
#         self.output_dens=nn.Dense(num_hiddens,activation="sigmoid")
#     def transpose(self,X,batch_size):
#
#         X=X.reshape(batch_size,-1,self.num_heads,10)
#         X=nd.transpose(X,axes=(0,2,1,3))
#         return X
#     def transpose_output(self,X,num_heads):
#         X=X.reshape(-1,num_heads,X.shape[1],X.shape[2])
#         X=nd.transpose(X,axes=(0,2,1,3))
#         return X.reshape(X.shape[0],X.shape[1],-1)
#     def forward(self,  queries,keys,values,valid_lens):
#         batch_size=queries.shape[0]
#         queries=self.transpose(self.query_dense(queries),batch_size)
#         keys=self.transpose(self.key_dense(keys),batch_size)
#         values=self.transpose(self.value_dense(values),batch_size)
#         output=scale_dot_product_attention(queries,keys,values,valid_lens)
#         output_concat=self.transpose_output(output,self.num_heads)
#         return self.output_dens(output_concat)
# if __name__ == '__main__':
#     num_hiddens,num_heads=100,5
#     attention=MultHeadAttention(num_hiddens,num_heads)
#     attention.initialize()
#     batch_size,num_queries,num_kvparis,valid_lens=2,4,6,nd.array([3,2])
#     queris=nd.random.uniform(shape=(batch_size,num_queries,num_hiddens))
#     values=nd.random.uniform(shape=(batch_size,num_kvparis,num_hiddens))
#     output=attention(queris,values,values,valid_lens)
#     print(output)