PyTorch 实战：从 0 开始搭建 Transformer

最新推荐文章于 2025-10-18 11:44:01 发布

原创

最新推荐文章于 2025-10-18 11:44:01 发布 · 1.6k 阅读

12 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #transformer #人工智能 #机器学习 #深度学习 #python #算法

导入必要的库

python

import math
import torch
import torch.nn as nn
from LabmL_helpers.module import Module
from labml_n.utils import clone_module_List
from typing import Optional, List
from torch.utils.data import DataLoader, TensorDataset
from torch import optim
import torch.nn.functional as F

Transformer 模型概述
Transformer 是一种序列到序列的模型，通过自注意力机制并行处理整个序列，能同时考虑序列中的所有元素，并学习上下文之间的关系。其架构包括编码器和解码器部分，每部分都由多个相同的层组成，这些层包含自注意力机制、前馈神经网络，以及归一化和 Dropout 步骤。
核心公式
- 自注意力计算：Attention(Q,K,V)=softmax(dkQKT)V，其中，Q、K、V分别是查询（Query）、键（Key）和值（Value）矩阵，dk是键的维度。
- 多头注意力：将输入分割为多个头，分别计算注意力，然后将结果拼接起来。
- 位置编码：由于 Transformer 不使用循环结构，因此引入位置编码来保留序列中的位置信息。
自注意力机制
- 核心原理：计算句子在编码过程中每个位置上的注意力权重，然后以权重和的方式来计算整个句子的隐含向量表示。公式中，首先将 query 与 key 的转置做点积，然后将结果除以dk ，再进行 softmax 计算，最后将结果与 value 做矩阵乘法得到 output。除以dk是为了防止QKT过大导致 softmax 计算溢出，且可使QKT结果满足均值为 0，方差 1 的分布。QKT计算本质上是余弦相似度，可表示两个向量在方向上的相似度。
- 实现

python

import numpy as np
from math import sqrt
import torch
from torch import nn


class Self_Attention(nn.Module):
    # input : batch_size * seq_len * input_dim
    # q : batch_size * input_dim * dim_k
    #