从零构建 Transformer 模型
从零实现 Transformer 是理解其架构的最佳方式。本文完整实现一个标准的 Encoder-Decoder Transformer。
完整实现
python
class Transformer(nn.Module):
def __init__(self, src_vocab, tgt_vocab, d_model=512, n_heads=8,
n_layers=6, d_ff=2048, dropout=0.1):
super().__init__()
self.encoder = nn.ModuleList([
EncoderLayer(d_model, n_heads, d_ff, dropout)
for _ in range(n_layers)
])
self.decoder = nn.ModuleList([
DecoderLayer(d_model, n_heads, d_ff, dropout)
for _ in range(n_layers)
])
self.src_embed = nn.Embedding(src_vocab, d_model)
self.tgt_embed = nn.Embedding(tgt_vocab, d_model)
self.pos_enc = PositionalEncoding(d_model, dropout)
self.proj = nn.Linear(d_model, tgt_vocab)
def forward(self, src, tgt, src_mask=None, tgt_mask=None):
src = self.pos_enc(self.src_embed(src))
tgt = self.pos_enc(self.tgt_embed(tgt))
for layer in self.encoder:
src = layer(src, src_mask)
for layer in self.decoder:
tgt = layer(tgt, src, src_mask, tgt_mask)
return self.proj(tgt)位置编码的选择
原始 Transformer 使用正弦位置编码。现代实现更常用可学习位置编码或 RoPE，后者支持外推到更长序列。