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

import torch
import torch.nn as nn
import math

class MultiHeadSelfAttention(nn.Module):
    def __init__(self,nums_head=8,hidden_dim=512):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.nums_head = nums_head
        self.head_dim = hidden_dim//nums_head
        self.q_proj = nn.Linear(hidden_dim,hidden_dim)
        self.k_proj = nn.Linear(hidden_dim,hidden_dim)
        self.v_proj = nn.Linear(hidden_dim,hidden_dim)
        self.o_proj = nn.Linear(hidden_dim,hidden_dim)
        self.dropout=nn.Dropout(0.1)

    def forward(self,x,attn_mask=None):
        batch_size,seq_len,_ = x.size()
        q=self.q_proj(x)
        k=self.k_proj(x)
        v=self.v_proj(x)
        # split
        q_state = q.view(batch_size,seq_len,self.nums_head,self.head_dim).transpose(1,2)
        k_state = k.view(batch_size,seq_len,self.nums_head,self.head_dim).transpose(1,2)
        v_state = v.view(batch_size,seq_len,self.nums_head,self.head_dim).transpose(1,2)
        
        # q@k^t/sqrt(d)
        attn_score = q_state@k_state.transpose(-1,-2)/math.sqrt(self.head_dim)

        if attn_mask is not None:
            attn_score = attn_score.masked_fill(attn_mask ==0,-1e9)
        attn_weight = torch.softmax(attn_score,dim=-1)
        attn_weight = self.dropout(attn_weight)
        output_mid = attn_weight@v_state
        # concat
        output_mid = output_mid.transpose(-1,-2).contiguous()
        output_mid = output_mid.view(batch_size,seq_len,-1)
        output= self.o_proj(output_mid)
        return output

x= torch.rand(3,2,512)
attn_mask = torch.tensor([
    [0,1],
    [0,0],
    [1,0],
]).unsqueeze(1).unsqueeze(2).expand(3,8,2,2)

net = MultiHeadSelfAttention(8,512)
net(x,attn_mask).shape