11import math
12
13import torch
14import torch.nn as nn
15from labml_helpers.module import Module
16
17from labml_nn.utils import clone_module_list
18from .feed_forward import FeedForward
19from .mha import MultiHeadAttention
20from .positional_encoding import get_positional_encoding
23class EmbeddingsWithPositionalEncoding(Module):
30 def __init__(self, d_model: int, n_vocab: int, max_len: int = 5000):
31 super().__init__()
32 self.linear = nn.Embedding(n_vocab, d_model)
33 self.d_model = d_model
34 self.register_buffer('positional_encodings', get_positional_encoding(d_model, max_len))
36 def forward(self, x: torch.Tensor):
37 pe = self.positional_encodings[:x.shape[0]].requires_grad_(False)
38 return self.linear(x) * math.sqrt(self.d_model) + pe
41class EmbeddingsWithLearnedPositionalEncoding(Module):
48 def __init__(self, d_model: int, n_vocab: int, max_len: int = 5000):
49 super().__init__()
50 self.linear = nn.Embedding(n_vocab, d_model)
51 self.d_model = d_model
52 self.positional_encodings = nn.Parameter(torch.zeros(max_len, 1, d_model), requires_grad=True)
54 def forward(self, x: torch.Tensor):
55 pe = self.positional_encodings[:x.shape[0]]
56 return self.linear(x) * math.sqrt(self.d_model) + pe
This can act as an encoder layer or a decoder layer.
🗒 Some implementations, including the paper seem to have differences in where the layer-normalization is done. Here we do a layer normalization before attention and feed-forward networks, and add the original residual vectors. Alternative is to do a layer normalization after adding the residuals. But we found this to be less stable when training. We found a detailed discussion about this in the paper On Layer Normalization in the Transformer Architecture.
59class TransformerLayer(Module):
d_model
is the token embedding size self_attn
is the self attention module src_attn
is the source attention module (when this is used in a decoder) feed_forward
is the feed forward module dropout_prob
is the probability of dropping out after self attention and FFN77 def __init__(self, *,
78 d_model: int,
79 self_attn: MultiHeadAttention,
80 src_attn: MultiHeadAttention = None,
81 feed_forward: FeedForward,
82 dropout_prob: float):
90 super().__init__()
91 self.size = d_model
92 self.self_attn = self_attn
93 self.src_attn = src_attn
94 self.feed_forward = feed_forward
95 self.dropout = nn.Dropout(dropout_prob)
96 self.norm_self_attn = nn.LayerNorm([d_model])
97 if self.src_attn is not None:
98 self.norm_src_attn = nn.LayerNorm([d_model])
99 self.norm_ff = nn.LayerNorm([d_model])
Whether to save input to the feed forward layer
101 self.is_save_ff_input = False
103 def forward(self, *,
104 x: torch.Tensor,
105 mask: torch.Tensor,
106 src: torch.Tensor = None,
107 src_mask: torch.Tensor = None):
Normalize the vectors before doing self attention
109 z = self.norm_self_attn(x)
Run through self attention, i.e. keys and values are from self
111 self_attn = self.self_attn(query=z, key=z, value=z, mask=mask)
Add the self attention results
113 x = x + self.dropout(self_attn)
If a source is provided, get results from attention to source. This is when you have a decoder layer that pays attention to encoder outputs
118 if src is not None:
Normalize vectors
120 z = self.norm_src_attn(x)
Attention to source. i.e. keys and values are from source
122 attn_src = self.src_attn(query=z, key=src, value=src, mask=src_mask)
Add the source attention results
124 x = x + self.dropout(attn_src)
Normalize for feed-forward
127 z = self.norm_ff(x)
Save the input to the feed forward layer if specified
129 if self.is_save_ff_input:
130 self.ff_input = z.clone()
Pass through the feed-forward network
132 ff = self.feed_forward(z)
Add the feed-forward results back
134 x = x + self.dropout(ff)
135
136 return x
139class Encoder(Module):
146 def __init__(self, layer: TransformerLayer, n_layers: int):
147 super().__init__()
Make copies of the transformer layer
149 self.layers = clone_module_list(layer, n_layers)
Final normalization layer
151 self.norm = nn.LayerNorm([layer.size])
153 def forward(self, x: torch.Tensor, mask: torch.Tensor):
Run through each transformer layer
155 for layer in self.layers:
156 x = layer(x=x, mask=mask)
Finally, normalize the vectors
158 return self.norm(x)
161class Decoder(Module):
168 def __init__(self, layer: TransformerLayer, n_layers: int):
169 super().__init__()
Make copies of the transformer layer
171 self.layers = clone_module_list(layer, n_layers)
Final normalization layer
173 self.norm = nn.LayerNorm([layer.size])
175 def forward(self, x: torch.Tensor, memory: torch.Tensor, src_mask: torch.Tensor, tgt_mask: torch.Tensor):
Run through each transformer layer
177 for layer in self.layers:
178 x = layer(x=x, mask=tgt_mask, src=memory, src_mask=src_mask)
Finally, normalize the vectors
180 return self.norm(x)
This predicts the tokens and gives the lof softmax of those. You don't need this if you are using nn.CrossEntropyLoss
.
183class Generator(Module):
193 def __init__(self, n_vocab: int, d_model: int):
194 super().__init__()
195 self.projection = nn.Linear(d_model, n_vocab)
197 def forward(self, x):
198 return self.projection(x)
201class EncoderDecoder(Module):
208 def __init__(self, encoder: Encoder, decoder: Decoder, src_embed: Module, tgt_embed: Module, generator: Module):
209 super().__init__()
210 self.encoder = encoder
211 self.decoder = decoder
212 self.src_embed = src_embed
213 self.tgt_embed = tgt_embed
214 self.generator = generator
This was important from their code. Initialize parameters with Glorot / fan_avg.
218 for p in self.parameters():
219 if p.dim() > 1:
220 nn.init.xavier_uniform_(p)
222 def forward(self, src: torch.Tensor, tgt: torch.Tensor, src_mask: torch.Tensor, tgt_mask: torch.Tensor):
Run the source through encoder
224 enc = self.encode(src, src_mask)
Run encodings and targets through decoder
226 return self.decode(enc, src_mask, tgt, tgt_mask)
228 def encode(self, src: torch.Tensor, src_mask: torch.Tensor):
229 return self.encoder(self.src_embed(src), src_mask)
231 def decode(self, memory: torch.Tensor, src_mask: torch.Tensor, tgt: torch.Tensor, tgt_mask: torch.Tensor):
232 return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)