| 30 | return src |
| 31 | |
| 32 | class Attention(nn.Module): |
| 33 | # calculate multi-head attention |
| 34 | def __init__(self, d_query, d_key, n_head, d_hidden, d_out, p_drop=0.1): |
| 35 | super(Attention, self).__init__() |
| 36 | self.h = n_head |
| 37 | self.dim = d_hidden |
| 38 | # |
| 39 | self.to_q = nn.Linear(d_query, n_head*d_hidden, bias=False) |
| 40 | self.to_k = nn.Linear(d_key, n_head*d_hidden, bias=False) |
| 41 | self.to_v = nn.Linear(d_key, n_head*d_hidden, bias=False) |
| 42 | # |
| 43 | self.to_out = nn.Linear(n_head*d_hidden, d_out) |
| 44 | self.scaling = 1/math.sqrt(d_hidden) |
| 45 | # |
| 46 | # initialize all parameters properly |
| 47 | self.reset_parameter() |
| 48 | |
| 49 | def reset_parameter(self): |
| 50 | # query/key/value projection: Glorot uniform / Xavier uniform |
| 51 | nn.init.xavier_uniform_(self.to_q.weight) |
| 52 | nn.init.xavier_uniform_(self.to_k.weight) |
| 53 | nn.init.xavier_uniform_(self.to_v.weight) |
| 54 | |
| 55 | # to_out: right before residual connection: zero initialize -- to make it sure residual operation is same to the Identity at the begining |
| 56 | nn.init.zeros_(self.to_out.weight) |
| 57 | nn.init.zeros_(self.to_out.bias) |
| 58 | |
| 59 | def forward(self, query, key, value): |
| 60 | B, Q = query.shape[:2] |
| 61 | B, K = key.shape[:2] |
| 62 | # |
| 63 | query = self.to_q(query).reshape(B, Q, self.h, self.dim) |
| 64 | key = self.to_k(key).reshape(B, K, self.h, self.dim) |
| 65 | value = self.to_v(value).reshape(B, K, self.h, self.dim) |
| 66 | # |
| 67 | query = query * self.scaling |
| 68 | attn = einsum('bqhd,bkhd->bhqk', query, key) |
| 69 | attn = F.softmax(attn, dim=-1) |
| 70 | # |
| 71 | out = einsum('bhqk,bkhd->bqhd', attn, value) |
| 72 | out = out.reshape(B, Q, self.h*self.dim) |
| 73 | # |
| 74 | out = self.to_out(out) |
| 75 | |
| 76 | return out |
| 77 | |
| 78 | class AttentionWithBias(nn.Module): |
| 79 | def __init__(self, d_in=256, d_bias=128, n_head=8, d_hidden=32): |