Transformer 代码实现#

本篇从零实现一个完整的 Transformer，涵盖前四章的所有组件。每个模块都配有注释，对应前文讲解的概念。

1. 实现路线图#

2. 完整代码#

1
import subprocess
2
subprocess.check_call(["pip", "install", "torch", "numpy"])
3

4
import torch
5
import torch.nn as nn
6
import torch.nn.functional as F
7
import numpy as np
8
import math
9

10
# ============================================================
11
# ① 位置编码 (对应 02_位置编码.md)
12
# ============================================================
13
class PositionalEncoding(nn.Module):
14
    """正弦位置编码，注入序列顺序信息"""
15

16
    def __init__(self, d_model, max_len=5000, dropout=0.1):
17
        super().__init__()
18
        self.dropout = nn.Dropout(dropout)
19

20
        pe = torch.zeros(max_len, d_model)             # (max_len, d_model)
21
        position = torch.arange(max_len).unsqueeze(1).float()  # (max_len, 1)
22
        # div_term: 10000^(2i/d_model)，用 exp+log 计算避免大数
23
        div_term = torch.exp(
24
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
25
        )
26
        pe[:, 0::2] = torch.sin(position * div_term)   # 偶数维度
27
        pe[:, 1::2] = torch.cos(position * div_term)   # 奇数维度
28
        pe = pe.unsqueeze(0)                            # (1, max_len, d_model)
29
        self.register_buffer('pe', pe)                  # 不参与梯度更新
30

31
    def forward(self, x):
32
        # x: (batch, seq_len, d_model)
33
        x = x + self.pe[:, :x.size(1)]
34
        return self.dropout(x)
35

36

37
# ============================================================
38
# ② Token Embedding (对应 01_Token Embedding.md)
39
# ============================================================
40
class TokenEmbedding(nn.Module):
41
    """Token Embedding + 缩放"""
42

43
    def __init__(self, vocab_size, d_model):
44
        super().__init__()
45
        self.embedding = nn.Embedding(vocab_size, d_model)
46
        self.d_model = d_model
47

48
    def forward(self, x):
49
        # 乘以 sqrt(d_model) 防止位置编码淹没语义信息
50
        return self.embedding(x) * math.sqrt(self.d_model)
51

52

53
# ============================================================
54
# ③ 多头注意力 (对应 03_多头注意力.md)
55
# ============================================================
56
class MultiHeadAttention(nn.Module):
57
    """Multi-Head Attention: 从多个角度并行关注"""
58

59
    def __init__(self, d_model, n_heads, dropout=0.1):
60
        super().__init__()
61
        assert d_model % n_heads == 0, "d_model 必须能被 n_heads 整除"
62

63
        self.d_model = d_model
64
        self.n_heads = n_heads
65
        self.d_k = d_model // n_heads  # 每个头的维度
66

67
        # 一次性投影 Q, K, V（高效实现）
68
        self.W_Q = nn.Linear(d_model, d_model)
69
        self.W_K = nn.Linear(d_model, d_model)
70
        self.W_V = nn.Linear(d_model, d_model)
71
        self.W_O = nn.Linear(d_model, d_model)  # 输出投影
72

73
        self.dropout = nn.Dropout(dropout)
74

75
    def forward(self, Q, K, V, mask=None):
76
        batch_size = Q.size(0)
77

78
        # 线性投影后拆分成多头: (batch, seq, d_model) → (batch, n_heads, seq, d_k)
79
        Q = self.W_Q(Q).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
80
        K = self.W_K(K).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
81
        V = self.W_V(V).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
82

83
        # Scaled Dot-Product Attention
84
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
85

86
        if mask is not None:
87
            scores = scores.masked_fill(mask == 0, -1e9)
88

89
        attn_weights = self.dropout(F.softmax(scores, dim=-1))
90
        attn_output = torch.matmul(attn_weights, V)
91

92
        # 拼接所有头: (batch, n_heads, seq, d_k) → (batch, seq, d_model)
93
        attn_output = attn_output.transpose(1, 2).contiguous().view(
94
            batch_size, -1, self.d_model
95
        )
96
        return self.W_O(attn_output)
97

98

99
# ============================================================
100
# ④ 前馈网络 (对应 01_Encoder Block.md §4)
101
# ============================================================
102
class FeedForward(nn.Module):
103
    """Position-wise Feed-Forward Network: 升维 → 激活 → 降维"""
104

105
    def __init__(self, d_model, d_ff, dropout=0.1):
106
        super().__init__()
107
        self.fc1 = nn.Linear(d_model, d_ff)
108
        self.fc2 = nn.Linear(d_ff, d_model)
109
        self.dropout = nn.Dropout(dropout)
110

111
    def forward(self, x):
112
        return self.fc2(self.dropout(F.relu(self.fc1(x))))
113

114

115
# ============================================================
116
# ⑤ Encoder Block (对应 01_Encoder Block.md)
117
# ============================================================
118
class EncoderBlock(nn.Module):
119
    """编码器块: Self-Attention → Add&Norm → FFN → Add&Norm"""
120

121
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
122
        super().__init__()
123
        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)
124
        self.ffn = FeedForward(d_model, d_ff, dropout)
125
        self.norm1 = nn.LayerNorm(d_model)
126
        self.norm2 = nn.LayerNorm(d_model)
127
        self.dropout1 = nn.Dropout(dropout)
128
        self.dropout2 = nn.Dropout(dropout)
129

130
    def forward(self, x, src_mask=None):
131
        # 子层 1: 多头自注意力 + 残差 + LayerNorm
132
        attn_out = self.self_attn(x, x, x, src_mask)
133
        x = self.norm1(x + self.dropout1(attn_out))
134
        # 子层 2: FFN + 残差 + LayerNorm
135
        ffn_out = self.ffn(x)
136
        x = self.norm2(x + self.dropout2(ffn_out))
137
        return x
138

139

140
# ============================================================
141
# ⑥ Decoder Block (对应 02_Masked Self Attention.md)
142
# ============================================================
143
class DecoderBlock(nn.Module):
144
    """解码器块: Masked Self-Attn → Cross-Attn → FFN，各带残差+Norm"""
145

146
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
147
        super().__init__()
148
        self.masked_attn = MultiHeadAttention(d_model, n_heads, dropout)
149
        self.cross_attn = MultiHeadAttention(d_model, n_heads, dropout)
150
        self.ffn = FeedForward(d_model, d_ff, dropout)
151
        self.norm1 = nn.LayerNorm(d_model)
152
        self.norm2 = nn.LayerNorm(d_model)
153
        self.norm3 = nn.LayerNorm(d_model)
154
        self.dropout1 = nn.Dropout(dropout)
155
        self.dropout2 = nn.Dropout(dropout)
156
        self.dropout3 = nn.Dropout(dropout)
157

158
    def forward(self, x, enc_output, tgt_mask=None, src_mask=None):
159
        # 子层 1: Masked Self-Attention（因果掩码，防止看到未来）
160
        attn_out = self.masked_attn(x, x, x, tgt_mask)
161
        x = self.norm1(x + self.dropout1(attn_out))
162
        # 子层 2: Cross-Attention（Q 来自解码器，K/V 来自编码器）
163
        cross_out = self.cross_attn(x, enc_output, enc_output, src_mask)
164
        x = self.norm2(x + self.dropout2(cross_out))
165
        # 子层 3: FFN
166
        ffn_out = self.ffn(x)
167
        x = self.norm3(x + self.dropout3(ffn_out))
168
        return x
169

170

171
# ============================================================
172
# ⑦ 完整 Transformer (对应 02_Transformer整体架构.md)
173
# ============================================================
174
class Transformer(nn.Module):
175
    """完整的 Encoder-Decoder Transformer"""
176

177
    def __init__(self, src_vocab, tgt_vocab, d_model=512, n_heads=8,
178
                 n_layers=6, d_ff=2048, max_len=5000, dropout=0.1):
179
        super().__init__()
180

181
        # Embedding + 位置编码
182
        self.src_embed = TokenEmbedding(src_vocab, d_model)
183
        self.tgt_embed = TokenEmbedding(tgt_vocab, d_model)
184
        self.pos_enc = PositionalEncoding(d_model, max_len, dropout)
185

186
        # 编码器和解码器堆叠 N 层
187
        self.encoder_layers = nn.ModuleList([
188
            EncoderBlock(d_model, n_heads, d_ff, dropout) for _ in range(n_layers)
189
        ])
190
        self.decoder_layers = nn.ModuleList([
191
            DecoderBlock(d_model, n_heads, d_ff, dropout) for _ in range(n_layers)
192
        ])
193

194
        # 输出投影层 (对应 03_终端输出.md)
195
        self.output_proj = nn.Linear(d_model, tgt_vocab)
196

197
    def encode(self, src, src_mask=None):
198
        """编码器前向传播"""
199
        x = self.pos_enc(self.src_embed(src))
200
        for layer in self.encoder_layers:
201
            x = layer(x, src_mask)
202
        return x
203

204
    def decode(self, tgt, enc_output, tgt_mask=None, src_mask=None):
205
        """解码器前向传播"""
206
        x = self.pos_enc(self.tgt_embed(tgt))
207
        for layer in self.decoder_layers:
208
            x = layer(x, enc_output, tgt_mask, src_mask)
209
        return x
210

211
    def forward(self, src, tgt, src_mask=None, tgt_mask=None):
212
        enc_output = self.encode(src, src_mask)
213
        dec_output = self.decode(tgt, enc_output, tgt_mask, src_mask)
214
        logits = self.output_proj(dec_output)  # (batch, tgt_len, tgt_vocab)
215
        return logits
216

217

218
# ============================================================
219
# 工具函数：生成掩码
220
# ============================================================
221
def generate_causal_mask(size):
222
    """生成因果掩码（下三角矩阵）"""
223
    mask = torch.tril(torch.ones(size, size)).unsqueeze(0).unsqueeze(0)
224
    return mask  # (1, 1, size, size)
225

226

227
# ============================================================
228
# ⑧ 测试：前向传播验证
229
# ============================================================
230
if __name__ == "__main__":
231
    # 超参数（缩小版，方便测试）
232
    src_vocab = 1000
233
    tgt_vocab = 1000
234
    d_model = 64
235
    n_heads = 4
236
    n_layers = 2
237
    d_ff = 256
238
    batch_size = 2
239
    src_len = 10
240
    tgt_len = 8
241

242
    # 创建模型
243
    model = Transformer(src_vocab, tgt_vocab, d_model, n_heads, n_layers, d_ff)
244

245
    # 随机输入
246
    src = torch.randint(0, src_vocab, (batch_size, src_len))
247
    tgt = torch.randint(0, tgt_vocab, (batch_size, tgt_len))
248
    tgt_mask = generate_causal_mask(tgt_len)
249

250
    # 前向传播
251
    logits = model(src, tgt, tgt_mask=tgt_mask)
252

253
    print(f"源序列形状:   {src.shape}")      # (2, 10)
254
    print(f"目标序列形状: {tgt.shape}")       # (2, 8)
255
    print(f"输出 logits:  {logits.shape}")   # (2, 8, 1000)
256
    print(f"模型参数量:   {sum(p.numel() for p in model.parameters()):,}")
257

258
    # 验证输出是有效的概率分布
259
    probs = F.softmax(logits[0, 0], dim=-1)
260
    print(f"\n第一个样本第一个位置的概率和: {probs.sum().item():.4f}")
261
    print(f"Top-5 预测: {torch.topk(probs, 5).indices.tolist()}")

3. 代码结构与前文对应#

1
Transformer
2
├── TokenEmbedding        ← 02 章 01_Token Embedding.md
3
├── PositionalEncoding     ← 02 章 02_位置编码.md
4
├── EncoderBlock ×N
5
│   ├── MultiHeadAttention ← 03 章 03_多头注意力.md
6
│   │   └── ScaledDotProduct ← 03 章 02_Self Attention计算.md
7
│   ├── Add & LayerNorm    ← 04 章 01_Encoder Block.md §2-3
8
│   └── FeedForward        ← 04 章 01_Encoder Block.md §4
9
├── DecoderBlock ×N
10
│   ├── Masked MHA         ← 04 章 02_Masked Self Attention.md
11
│   ├── Cross MHA          ← 04 章 02_Masked Self Attention.md §4
12
│   ├── Add & LayerNorm
13
│   └── FeedForward
14
└── output_proj (Linear)   ← 04 章 03_终端输出.md

4. 训练循环示例#

1
# 接上面的代码，展示训练循环核心逻辑（伪代码）
2

3
# 优化器 + 损失函数
4
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, betas=(0.9, 0.98), eps=1e-9)
5
criterion = nn.CrossEntropyLoss(ignore_index=0, label_smoothing=0.1)
6

7
# 单步训练
8
model.train()
9
optimizer.zero_grad()
10

11
# 前向传播
12
logits = model(src, tgt[:, :-1], tgt_mask=generate_causal_mask(tgt_len - 1))
13
# logits: (batch, tgt_len-1, vocab_size)
14

15
# 计算损失（展平后计算交叉熵）
16
target = tgt[:, 1:]  # 右移一位作为目标
17
loss = criterion(logits.reshape(-1, tgt_vocab), target.reshape(-1))
18

19
# 反向传播
20
loss.backward()
21

22
# 梯度裁剪
23
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
24

25
# 参数更新
26
optimizer.step()
27

28
print(f"Loss: {loss.item():.4f}")

从这个基础出发
这份代码是”教学版”——结构清晰但未做工程优化。实际生产中还需要：

Flash Attention：融合 kernel，显存和速度大幅优化

混合精度训练 (FP16/BF16)：减少显存，加速计算

分布式训练 (DDP/FSDP)：多卡/多机并行

KV Cache：推理时缓存已计算的 K/V，避免重复计算

Transformer 代码实现#

1. 实现路线图#

2. 完整代码#

3. 代码结构与前文对应#

4. 训练循环示例#

相关笔记#