-
Notifications
You must be signed in to change notification settings - Fork 0
/
EncoderUnit.py
32 lines (24 loc) · 1.26 KB
/
EncoderUnit.py
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
import torch
import torch.nn as nn
from .SelfAttention import SelfAttention
class EncoderUnit(nn.Module):
def __init__(self, embed_size: int, heads: int, forward_expansion: int, dropout: float = 0.5):
super(EncoderUnit, self).__init__()
self.embed_size = embed_size
self.heads = heads
self.dropout = dropout
self.forward_expansion = forward_expansion
self.self_attention = SelfAttention(embed_size, heads) # dim: batch_size * q * embed_size
self.norm1 = nn.LayerNorm(embed_size)
self.feed_forward = nn.Sequential(
nn.Linear(embed_size, forward_expansion * embed_size),
nn.ReLU(),
nn.Linear(forward_expansion * embed_size, embed_size)
)
self.norm2 = nn.LayerNorm(embed_size)
self.dropout = nn.Dropout(dropout)
def forward(self, values: torch.Tensor, keys: torch.Tensor, queries: torch.Tensor, mask: torch.Tensor = None):
attn_scores = self.self_attention(values, keys, queries, mask)
normalised_score = self.dropout(self.norm1(attn_scores + queries))
feed_forward_score = self.feed_forward(normalised_score)
return self.dropout(self.norm2(feed_forward_score + normalised_score))