13import math
14
15import torch
16import torch.nn as nn
17
18from labml_nn.utils import clone_module_list
19from .feed_forward import FeedForward
20from .mha import MultiHeadAttention
21from .positional_encoding import get_positional_encoding
31 def __init__(self, d_model: int, n_vocab: int, max_len: int = 5000):
32 super().__init__()
33 self.linear = nn.Embedding(n_vocab, d_model)
34 self.d_model = d_model
35 self.register_buffer('positional_encodings', get_positional_encoding(d_model, max_len))
37 def forward(self, x: torch.Tensor):
38 pe = self.positional_encodings[:x.shape[0]].requires_grad_(False)
39 return self.linear(x) * math.sqrt(self.d_model) + pe
42class EmbeddingsWithLearnedPositionalEncoding(nn.Module):
49 def __init__(self, d_model: int, n_vocab: int, max_len: int = 5000):
50 super().__init__()
51 self.linear = nn.Embedding(n_vocab, d_model)
52 self.d_model = d_model
53 self.positional_encodings = nn.Parameter(torch.zeros(max_len, 1, d_model), requires_grad=True)
55 def forward(self, x: torch.Tensor):
56 pe = self.positional_encodings[:x.shape[0]]
57 return self.linear(x) * math.sqrt(self.d_model) + pe
它可以充当编码器层或解码器层。
🗒 包括论文在内的一些实现似乎在图层归一化的位置上存在差异。在这里,我们在注意力和前馈网络之前进行层归一化,并添加原始残差向量。另一种方法是在添加残差后进行图层归一化。但是我们发现在训练时这种情况不太稳定。我们在《变压器架构中的层规范化》一文中找到了对此的详细讨论。
60class TransformerLayer(nn.Module):
d_model
是令牌嵌入的大小self_attn
是自我关注模块src_attn
是源关注模块(当它在解码器中使用时)feed_forward
是前馈模块dropout_prob
是自我关注和 FFN 后退学的概率78 def __init__(self, *,
79 d_model: int,
80 self_attn: MultiHeadAttention,
81 src_attn: MultiHeadAttention = None,
82 feed_forward: FeedForward,
83 dropout_prob: float):
91 super().__init__()
92 self.size = d_model
93 self.self_attn = self_attn
94 self.src_attn = src_attn
95 self.feed_forward = feed_forward
96 self.dropout = nn.Dropout(dropout_prob)
97 self.norm_self_attn = nn.LayerNorm([d_model])
98 if self.src_attn is not None:
99 self.norm_src_attn = nn.LayerNorm([d_model])
100 self.norm_ff = nn.LayerNorm([d_model])
是否将输入保存到前馈层
102 self.is_save_ff_input = False
104 def forward(self, *,
105 x: torch.Tensor,
106 mask: torch.Tensor,
107 src: torch.Tensor = None,
108 src_mask: torch.Tensor = None):
在进行自我注意之前对向量进行归一化
110 z = self.norm_self_attn(x)
通过自我关注,即关键和价值来自自我
112 self_attn = self.self_attn(query=z, key=z, value=z, mask=mask)
添加自我关注的结果
114 x = x + self.dropout(self_attn)
如果提供了来源,则从关注源获取结果。这是当你有一个关注编码器输出的解码器层
时119 if src is not None:
归一化向量
121 z = self.norm_src_attn(x)
注意源。即键和值来自源
123 attn_src = self.src_attn(query=z, key=src, value=src, mask=src_mask)
添加来源关注结果
125 x = x + self.dropout(attn_src)
标准化以进行前馈
128 z = self.norm_ff(x)
如果已指定,则将输入保存到前馈图层
130 if self.is_save_ff_input:
131 self.ff_input = z.clone()
通过前馈网络
133 ff = self.feed_forward(z)
将前馈结果添加回来
135 x = x + self.dropout(ff)
136
137 return x
140class Encoder(nn.Module):
147 def __init__(self, layer: TransformerLayer, n_layers: int):
148 super().__init__()
制作变压器层的副本
150 self.layers = clone_module_list(layer, n_layers)
最终归一化层
152 self.norm = nn.LayerNorm([layer.size])
154 def forward(self, x: torch.Tensor, mask: torch.Tensor):
穿过每个变压器层
156 for layer in self.layers:
157 x = layer(x=x, mask=mask)
最后,对向量进行归一化
159 return self.norm(x)
162class Decoder(nn.Module):
169 def __init__(self, layer: TransformerLayer, n_layers: int):
170 super().__init__()
制作变压器层的副本
172 self.layers = clone_module_list(layer, n_layers)
最终归一化层
174 self.norm = nn.LayerNorm([layer.size])
176 def forward(self, x: torch.Tensor, memory: torch.Tensor, src_mask: torch.Tensor, tgt_mask: torch.Tensor):
穿过每个变压器层
178 for layer in self.layers:
179 x = layer(x=x, mask=tgt_mask, src=memory, src_mask=src_mask)
最后,对向量进行归一化
181 return self.norm(x)
184class Generator(nn.Module):
194 def __init__(self, n_vocab: int, d_model: int):
195 super().__init__()
196 self.projection = nn.Linear(d_model, n_vocab)
198 def forward(self, x):
199 return self.projection(x)
202class EncoderDecoder(nn.Module):
209 def __init__(self, encoder: Encoder, decoder: Decoder, src_embed: nn.Module, tgt_embed: nn.Module, generator: nn.Module):
210 super().__init__()
211 self.encoder = encoder
212 self.decoder = decoder
213 self.src_embed = src_embed
214 self.tgt_embed = tgt_embed
215 self.generator = generator
从他们的代码来看,这很重要。使用 Glorot/fan_avg 初始化参数。
219 for p in self.parameters():
220 if p.dim() > 1:
221 nn.init.xavier_uniform_(p)
223 def forward(self, src: torch.Tensor, tgt: torch.Tensor, src_mask: torch.Tensor, tgt_mask: torch.Tensor):
通过编码器运行源码
225 enc = self.encode(src, src_mask)
通过解码器运行编码和目标
227 return self.decode(enc, src_mask, tgt, tgt_mask)
229 def encode(self, src: torch.Tensor, src_mask: torch.Tensor):
230 return self.encoder(self.src_embed(src), src_mask)
232 def decode(self, memory: torch.Tensor, src_mask: torch.Tensor, tgt: torch.Tensor, tgt_mask: torch.Tensor):
233 return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)