File size: 3,521 Bytes
34393ef | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | import os
import sys
import torch
from torch import nn
import numpy as np
class self_attention(nn.Module):
def __init__(self,in_channel,n_head,d_k,d_v):
super(self_attention,self).__init__()
self.n_head=n_head
self.d_k = d_k
self.d_v = d_v
self.W_dict = nn.ModuleDict({"Wq" : nn.Linear(in_channel,n_head*d_k),
"Wk" : nn.Linear(in_channel,n_head*d_k),
"Wv" : nn.Linear(in_channel,n_head*d_v)})
def forward(self,X):
# some dimension
dk_sqrt = int(np.sqrt(self.d_k))
querys = self.W_dict['Wq'](X) # B*X_dim*dk
keys = self.W_dict['Wk'](X) # B* hs_out -> B*64
values = self.W_dict['Wv'](X) # B* hs_out -> B*128
sim_M = torch.bmm(querys,keys.transpose(1,2))/8 # B* X_dim * X_dim
attention = torch.softmax(sim_M,dim=-1)
# result
result = torch.bmm(attention,values).squeeze(2) # B*X_dim*1 -> B*X_dim
return result
class task_attention(nn.Module):
"""
Task Attention Layer, it contains a task specific mask `Wq` as the Query. Key is computed from the input
Arguments:
d_v : int , dimension of the network width input , i.e : (Batch_size, d_v)
d_k : int , dimension of the Query and Key.
"""
def __init__(self,d_v:int,d_k=64):
super(task_attention,self).__init__()
self.d_k = d_k
self.d_v = d_v
self.W_dict = nn.ModuleDict({"Wq" : nn.Linear(d_k,d_v),
"Wk" : nn.Linear(d_v,d_k)})
def forward(self,X):
# i.e X : (B ,128)
# if len(X.shape) == 2:
# X = X
# some dimension
dk_sqrt = int(np.sqrt(self.d_k))
keys = self.W_dict['Wk'](X).squeeze(2) # (B, 128) -> (B ,64)
# dot product similarity is used here , which is implemented by `nn.Linear`
sim_M =self.W_dict['Wq'](keys) /dk_sqrt # (B,128,64) * (B,64,128) -> (B, 128, 128)
attention = torch.softmax(sim_M,dim=-1)
# result
result = torch.mul(attention,X) # (B, 128) * (B, 128) -> B*128
return result
class task_attention(nn.Module):
"""
Task Attention Layer, it contains a task specific mask `Wq` as the Query. Key is computed from the input
Arguments:
d_v : int , dimension of the network width input , i.e : (Batch_size, d_v)
d_k : int , dimension of the Query and Key.
"""
def __init__(self,d_v:int,d_k=64):
super(task_attention,self).__init__()
self.d_k = d_k
self.d_v = d_v
self.W_dict = nn.ModuleDict({"Wq" : nn.Linear(d_v,d_k),
"Wk" : nn.Linear(d_v,d_k)})
def forward(self,X):
# i.e X : (B ,128)
# if len(X.shape) == 2:
# X = X
# some dimension
dk_sqrt = int(np.sqrt(self.d_k))
query = self.W_dict['Wq'](X).squeeze(2)
keys = self.W_dict['Wk'](X).squeeze(2) # (B, 128) -> (B ,64)
# dot product similarity is used here , which is implemented by `nn.Linear`
sim_M =self.W_dict['Wq'](keys) /dk_sqrt # (B,128,64) * (B,64,128) -> (B, 128, 128)
attention = torch.softmax(sim_M,dim=-1)
# result
result = torch.mul(attention,X) # (B, 128) * (B, 128) -> B*128
return result |