AI基础模型搭建实战：从零构建你的第一个Transformer模型

1次阅读

共计 3540 个字符，预计需要花费 9 分钟才能阅读完成。

第一次尝试搭建 AI 模型时，大多数初学者都会遇到这三个经典问题：

架构选择盲目 ：在 CNN、RNN 和 Transformer 之间纠结，不知道哪种结构适合自己的任务
训练资源浪费 ：好不容易跑起来的模型，却因为显存不足或参数设置不当导致训练效率低下
调试效率低下 ：模型效果不佳时，不知道从哪里开始排查问题

让我们先快速对比下主流神经网络结构的特性：

CNN（卷积神经网络）：
优势：局部特征提取能力强，适合图像等网格数据
局限：难以建模长距离依赖关系
RNN（循环神经网络）：
优势：适合处理序列数据（如文本、时间序列）
局限：梯度消失问题严重，并行计算能力差
Transformer：
自注意力机制可以同时捕获局部和全局特征
天然支持并行计算，训练效率更高
在 NLP 任务中表现出色（如 BERT、GPT 都基于 Transformer）

import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, n_heads=8):
        super().__init__()
        assert d_model % n_heads == 0, "d_model 必须能被 n_heads 整除"

        # 初始化参数
        self.d_model = d_model  # 模型维度
        self.n_heads = n_heads  # 注意力头数
        self.d_k = d_model // n_heads  # 每个头的维度

        # 线性变换层
        self.W_q = nn.Linear(d_model, d_model)  # Query 变换
        self.W_k = nn.Linear(d_model, d_model)  # Key 变换
        self.W_v = nn.Linear(d_model, d_model)  # Value 变换
        self.W_o = nn.Linear(d_model, d_model)  # 输出变换

    def forward(self, q, k, v, mask=None):
        # q/k/ v 形状: [batch_size, seq_len, d_model]
        batch_size = q.size(0)

        # 1. 线性变换并分头
        Q = self.W_q(q).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)  # [bs, n_heads, q_len, d_k]
        K = self.W_k(k).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)  # [bs, n_heads, k_len, d_k]
        V = self.W_v(v).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)  # [bs, n_heads, v_len, d_k]

        # 2. 计算注意力分数
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)  # [bs, n_heads, q_len, k_len]

        # 3. 可选 mask（如处理 padding）if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        # 4. 计算注意力权重
        attn_weights = torch.softmax(scores, dim=-1)

        # 5. 加权求和
        output = torch.matmul(attn_weights, V)  # [bs, n_heads, q_len, d_k]

        # 6. 合并多头输出
        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)  # [bs, q_len, d_model]

        # 7. 最终线性变换
        return self.W_o(output)

位置编码（Positional Encoding）是 Transformer 理解序列顺序的关键。我们可以用以下代码可视化位置编码的模式：

import matplotlib.pyplot as plt
import numpy as np

def get_positional_encoding(max_len, d_model):
    position = np.arange(max_len)[:, np.newaxis]
    div_term = np.exp(np.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))

    pe = np.zeros((max_len, d_model))
    pe[:, 0::2] = np.sin(position * div_term)  # 偶数位置用 sin
    pe[:, 1::2] = np.cos(position * div_term)  # 奇数位置用 cos
    return pe

pe = get_positional_encoding(100, 512)
plt.figure(figsize=(10, 5))
plt.imshow(pe.T, aspect='auto', cmap='viridis')
plt.colorbar()
plt.show()

你会看到类似海浪的图案，这说明不同位置有独特的编码模式，且相邻位置编码变化平滑。

当 GPU 显存不足以支持较大 batch size 时，可以使用梯度累积技术：

gradient_accumulation_steps = 4  # 累积 4 个 batch 的梯度

for epoch in range(epochs):
    model.train()
    optimizer.zero_grad()

    for i, (inputs, targets) in enumerate(train_loader):
        outputs = model(inputs)
        loss = criterion(outputs, targets)

        # 梯度累积
        loss = loss / gradient_accumulation_steps
        loss.backward()

        if (i+1) % gradient_accumulation_steps == 0:
            optimizer.step()
            optimizer.zero_grad()

使用 PyTorch 的 AMP（Automatic Mixed Precision）可以显著减少显存占用：

训练模式	显存占用 (GB)	训练速度 (iter/s)
FP32 (单精度)	12.7	3.2
FP16 (混合精度)	7.3 (-43%)	5.8 (+81%)

启用方法：

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

with autocast():
    outputs = model(inputs)
    loss = criterion(outputs, targets)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

Transformer 训练通常需要学习率 warmup，但如何设置持续时间呢？经验法则：