Daniel Voigt Godoy - Deep Learning with PyTorch Step-by-Step A Beginner’s Guide-leanpub

6. Decoder "Layer"
The decoder "layer" implements a list of three "sub-layers" (7), which are going to
be called with their corresponding operations:
Decoder "Layer"
1 class DecoderLayer(nn.Module):
2 def __init__(self, n_heads, d_model, ff_units, dropout=0.1):
3 super().__init__()
4 self.n_heads = n_heads
5 self.d_model = d_model
6 self.ff_units = ff_units
7 self.self_attn_heads = \
8 MultiHeadedAttention(n_heads, d_model, dropout)
9 self.cross_attn_heads = \
10 MultiHeadedAttention(n_heads, d_model, dropout)
11 self.ffn = nn.Sequential(nn.Linear(d_model, ff_units),
12 nn.ReLU(),
13 nn.Dropout(dropout),
14 nn.Linear(ff_units, d_model))
15 self.sublayers = nn.ModuleList(
16 [SubLayerWrapper(d_model, dropout) for _ in range(3)]
17 )
18
19 def init_keys(self, states):
20 self.cross_attn_heads.init_keys(states)
21
22 def forward(self, query, source_mask=None, target_mask=None):
23 # SubLayer 0 - Masked Self-Attention
24 att1 = self.sublayers[0](query, mask=target_mask, 1
25 sublayer=self.self_attn_heads,
26 is_self_attn=True)
27 # SubLayer 1 - Cross-Attention
28 att2 = self.sublayers[1](att1, mask=source_mask, 1
29 sublayer=self.cross_attn_heads)
30 # SubLayer 2 - FFN
31 out = self.sublayers[2](att2, sublayer=self.ffn) 1
32 return out
1 Calls "Sub-Layer" wrapper (7) three times (self-attention, cross-attention, and
feed-forward network)
870 | Chapter 10: Transform and Roll Out