transfomer网络构建

目录

整体架构

详细模块分析

Transformer 算法的演进

代码讲解

1. 导入必要的库

2. InputEmbeddings 类

3. PositionalEncoding 类

4. LayerNormalization 类

5. FeedForwardBlock 类

6. MultiHeadAttentionBlock 类

7. ResidualConnection 类

8. EncoderBlock 类

9. Encoder 类

10. DecoderBlock 类

11. Decoder 类

12. ProjectionLayer 类

13. Transformer 类

14. build_transformer 函数

完整代码实现transfomer

---

整体架构

项目的核心目标是构建一个基于 Transformer 架构的神经网络模型，用于实现英法（en - it）双语的机器翻译任务。

配置管理：借助配置文件来设定模型训练所需的各类参数，像批次大小、训练轮数、学习率等。

数据集处理：加载数据集，对数据进行分词处理，并且构建自定义的数据集类，从而适配模型的输入要求。

模型构建：依据 Transformer 架构构建模型，包含输入嵌入层、位置编码层、编码器和解码器等组件。

训练过程：对模型进行训练，保存模型参数，同时使用 TensorBoard 记录训练过程中的损失值。

详细模块分析

配置管理 (config.py)

定义和获取模型训练所需的配置参数

def get_config():
    return {
        "batch_size": 8,
        "num_epochs": 20,
        "lr": 10 ** -4,
        "seq_len": 350,
        "d_model": 512,
        "datasource": 'opus_books',
        "lang_src": "en",
        "lang_tgt": "it",
        "model_folder": "weights",
        "model_basename": "tmodel_",
        "tokenizer_file": "tokenizer_{0}.json",
        "experiment_name": "runs/tmodel"
    }

这些参数涵盖了训练的基本设置、数据集信息、模型保存路径等内容。

数据集处理 (dataset.py)

该模块的主要任务是构建自定义的数据集类 BilingualDataset

对原始文本进行分词处理

为输入序列添加特殊标记（如 [SOS]、[EOS]、[PAD]）

对序列进行填充，使其长度一致

生成编码器和解码器的掩码

class BilingualDataset(Dataset):
    def __init__(self, ds, tokenizer_src, tokenizer_tgt, src_lang, tgt_lang, seq_len):
        # 初始化数据集和分词器等
        ...    def __getitem__(self, index):
        # 获取数据并进行处理
        ...

模型构建 (model.py)

此模块按照 Transformer 架构构建模型，主要包含以下组件：

输入嵌入层 (InputEmbeddings)：把输入的词索引转换为词向量。

位置编码层 (PositionalEncoding)：为输入序列添加位置信息。

多头注意力机制 (MultiHeadAttentionBlock)：在不同的子空间中计算注意力分数。

前馈网络层 (FeedForwardBlock)：对多头注意力机制的输出进行非线性变换。

编码器和解码器 (Encoder 和 Decoder)：由多个编码器块和解码器块组成。

投影层 (ProjectionLayer)：把解码器的输出映射到词表大小的维度。

class Transformer(nn.Module):
    def __init__(self, encoder, decoder, src_embed, tgt_embed, src_pos, tgt_pos, projection_layer):
        # 初始化模型组件
        ...    def encode(self, src, src_mask):
        # 编码器前向传播
        ...    def decode(self, encoder_output, src_mask, tgt, tgt_mask):
        # 解码器前向传播
        ...    def project(self, x):
        # 投影层前向传播
        ...

训练过程 (train.py)

该模块负责模型的训练，具体步骤如下：

加载数据集和分词器。

构建模型和优化器。

定义损失函数。

分轮次和批次进行训练，同时记录损失值。

每个轮次结束后保存模型参数。

def train_model(config):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # 确保保存模型的文件夹存在
    ...
    train_dataloader, val_dataloader, tokenizer_src, tokenizer_tgt = get_ds(config)
    model = get_model(config, tokenizer_src.get_vocab_size(), tokenizer_tgt.get_vocab_size()).to(device)
    writer = SummaryWriter(config['experiment_name'])
    optimizer = torch.optim.Adam(model.parameters(), lr=config['lr'], eps=1e-9)
    loss_fn = nn.CrossEntropyLoss(ignore_index=tokenizer_src.token_to_id('[PAD]'), label_smoothing=0.1)
    # 分轮次和批次进行训练
    ...

Transformer 算法的演进

原始 Transformer

Transformer 最初在论文《Attention Is All You Need》中被提出，其核心创新点在于摒弃了传统的循环神经网络（RNN）和卷积神经网络（CNN），完全基于注意力机制构建模型。这种架构具有以下优势：

并行计算能力：能够并行处理输入序列，从而显著提高训练速度。

长序列处理能力：通过注意力机制，能够有效捕捉序列中的长距离依赖关系。

BERT（Bidirectional Encoder Representations from Transformers）

BERT 是基于 Transformer 编码器构建的预训练语言模型，其主要创新点在于：

双向训练：采用双向的注意力机制，能够同时考虑上下文信息，从而学习到更强大的语言表示。

预训练 - 微调（Pretraining - Fine-tuning）：先在大规模无监督数据上进行预训练，然后在特定任务上进行微调，大大提高了模型在各种自然语言处理任务上的性能。

GPT（Generative Pretrained Transformer）

GPT 是基于 Transformer 解码器构建的预训练语言模型，其主要特点是：

单向训练：采用单向的注意力机制，从左到右生成文本。

生成式任务：在文本生成任务上表现出色，如对话生成、文本摘要等。

代码讲解

1. 导入必要的库

import torch
import torch.nn as nn
import math

torch：PyTorch 深度学习框架的核心库，提供了张量操作和自动求导等功能。

torch.nn：包含了构建神经网络所需的各种模块和层。

math：Python 标准库中的数学模块，用于进行数学计算。

2. InputEmbeddings 类

class InputEmbeddings(nn.Module):
    # 词嵌入本质会训练一个矩阵, (vocab_size, d_model)
    def __init__(self, d_model: int, vocab_size: int):
        super().__init__()
        self.d_model = d_model
        self.vocab_size = vocab_size
        self.embedding = nn.Embedding(vocab_size, d_model)    def forward(self, x):
        # nn.Embedding 里面是有参数的, 使用的时候相当于 lookup table
        # 乘上权重来自于论文 3.4
        return self.embedding(x) * math.sqrt(self.d_model)

功能：将输入的词索引转换为对应的词向量。

初始化：d_model：词向量的维度。

vocab_size：词汇表的大小。

nn.Embedding(vocab_size, d_model)：创建一个嵌入层，将词索引映射到 d_model 维的向量空间。

前向传播：self.embedding(x)：通过嵌入层将输入的词索引转换为词向量。

* math.sqrt(self.d_model)：乘以 sqrt(d_model) 是为了缩放词向量，这是论文中的做法。

3. PositionalEncoding 类

class PositionalEncoding(nn.Module):
    # d_model 是每个位置的向量维度, seq_len 是总共多少个位置, dropout 防止过拟合
    def __init__(self, d_model: int, seq_len: int, dropout: float):
        super().__init__()
        self.d_model = d_model
        self.seq_len = seq_len
        self.dropout = nn.Dropout(dropout)        # 创建一个矩阵, 形状(seq_len, d_model)
        pe = torch.zeros(seq_len, d_model)
        # 创建一个向量 (seq_len, 1)
        position = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1)
        # 实现公式 3.5, 但是在 log space 下计算的
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)        pe = pe.unsqueeze(0)  # (1, seq_len, d_model)        # 不需要去训练的参数, 最终它也是会保存在模型文件中的, 可以把它注册到 buffer 中
        self.register_buffer('pe', pe)    def forward(self, x):
        # (batch, seq_len, d_model)
        x = x + (self.pe[:, :x.shape[1], :]).requires_grad_(False)
        return self.dropout(x)

功能：为输入的词向量添加位置信息，因为 Transformer 模型本身不具备捕捉序列顺序的能力。

初始化：d_model：词向量的维度。

seq_len：序列的最大长度。

dropout：Dropout 概率，用于防止过拟合。

pe：位置编码矩阵，形状为 (seq_len, d_model)。

position：位置向量，形状为 (seq_len, 1)。

div_term：用于计算位置编码的分母项。

self.register_buffer('pe', pe)：将位置编码矩阵注册为缓冲区，这样它不会被视为模型的可训练参数，但会随模型一起保存。

前向传播：x + (self.pe[:, :x.shape[1], :]).requires_grad_(False)：将位置编码矩阵添加到输入的词向量上。

self.dropout(x)：对添加了位置编码的词向量应用 Dropout。

4. LayerNormalization 类

class LayerNormalization(nn.Module):# layer norm 是对每个样本不同位置的向量求均值和方差,然后再进行归一化# features 是一个样本的特征数量def __init__(self, features: int, eps: float = 1e-6):super().__init__()
        self.eps = eps
        self.alpha = nn.Parameter(torch.ones(features))
        self.bias = nn.Parameter(torch.zeros(features))def forward(self, x):# x: (batch, seq_len, hidden_size) 每条文本的每个token由多个特征构成
        mean = x.mean(dim=-1, keepdim=True)
        std = x.std(dim=-1, keepdim=True)return self.alpha * (x - mean) / (std + self.eps) + self.bias

功能：对输入的张量进行层归一化操作，使得每个样本的特征向量具有零均值和单位方差。

初始化：features：输入张量的特征维度。

eps：一个小的常数，用于避免分母为零。

self.alpha：可训练的缩放因子，形状为 (features,)。

self.bias：可训练的偏移量，形状为 (features,)。

前向传播：mean = x.mean(dim=-1, keepdim=True)：计算每个样本的特征向量的均值。

std = x.std(dim=-1, keepdim=True)：计算每个样本的特征向量的标准差。

self.alpha * (x - mean) / (std + self.eps) + self.bias：进行层归一化操作，并应用缩放因子和偏移量。

5. FeedForwardBlock 类

class FeedForwardBlock(nn.Module):def __init__(self, d_model: int, d_ff: int, dropout: float):super().__init__()
        self.linear_1 = nn.Linear(d_model, d_ff)
        self.dropout = nn.Dropout(p=dropout)
        self.linear_2 = nn.Linear(d_ff, d_model)def forward(self, x):# 实现的就是论文中 3.3 FFN(x)return self.linear_2(self.dropout(torch.relu(self.linear_1(x))))

功能：实现一个前馈神经网络块，包含两个线性层和一个 ReLU 激活函数。

初始化：d_model：输入和输出的特征维度。

d_ff：中间层的特征维度。

dropout：Dropout 概率，用于防止过拟合。

self.linear_1：第一个线性层，将输入从 d_model 维映射到 d_ff 维。

self.linear_2：第二个线性层，将中间层的输出从 d_ff 维映射回 d_model 维。

前向传播：torch.relu(self.linear_1(x))：对输入应用第一个线性层和 ReLU 激活函数。

self.dropout(...)：对激活函数的输出应用 Dropout。

self.linear_2(...)：对 Dropout 的输出应用第二个线性层。

6. MultiHeadAttentionBlock 类

class MultiHeadAttentionBlock(nn.Module):
    def __init__(self, d_model: int, h: int, dropout: float):
        super().__init__()        self.d_model = d_model
        self.h = h
        assert d_model % h == 0, "d_model 必须可以整除 h"
        # 每个 head 对应的维度是多大
        self.d_k = d_model // h        self.w_q = nn.Linear(d_model, d_model, bias=False)
        self.w_k = nn.Linear(d_model, d_model, bias=False)
        self.w_v = nn.Linear(d_model, d_model, bias=False)
        self.w_o = nn.Linear(d_model, d_model, bias=False)        self.dropout = nn.Dropout(dropout)    @staticmethod
    def attention(query, key, value, mask, dropout: nn.Dropout):
        d_k = query.shape[-1]        attention_scores = (query @ key.transpose(-2, -1)) / math.sqrt(d_k)
        if mask is not None:
            # 将 mask 矩阵中值为 0的位置,对应在 attention_scores 里面的值用 -1e9 替换
            attention_scores.masked_fill_(mask == 0, -1e9)
        attention_scores = attention_scores.softmax(dim=-1)  # (Batch, h, seq_len, seq_len)
        if dropout is not None:
            attention_scores = dropout(attention_scores)        return (attention_scores @ value), attention_scores    def forward(self, q, k, v, mask):
        query = self.w_q(q)  # Q
        key = self.w_k(k)  # K
        value = self.w_v(v)  # V        # 多头注意力机制 切割 Q K V 为了给不同的 heads
        # (Batch, Seq_Len, d_model) --> (Batch, Seq_Len, h, d_k) --> (Batch, h, Seq_Len, d_k)
        query = query.view(query.shape[0], query.shape[1], self.h, self.d_k).transpose(1, 2)
        key = key.view(key.shape[0], key.shape[1], self.h, self.d_k).transpose(1, 2)
        value = value.view(value.shape[0], value.shape[1], self.h, self.d_k).transpose(1, 2)        x, self.attention_scores = MultiHeadAttentionBlock.attention(query, key, value, mask, self.dropout)
        # (Batch, h, seq_len, d_k) --> (Batch, Seq_Len, h, d_k) --> (Batch, Seq_Len, d_model)
        x = x.transpose(1, 2).contiguous().view(x.shape[0], -1, self.h * self.d_k)        return self.w_o(x)

功能：实现多头注意力机制，通过多个头并行计算注意力分数，然后将结果拼接并线性变换。

初始化：d_model：输入和输出的特征维度。

h：头的数量。

self.d_k：每个头的特征维度，d_k = d_model // h。

self.w_q、self.w_k、self.w_v：分别用于生成查询（Q）、键（K）和值（V）的线性层。

self.w_o：用于将多头注意力的输出进行线性变换的线性层。

self.dropout：Dropout 概率，用于防止过拟合。

attention 静态方法：计算注意力分数：attention_scores = (query @ key.transpose(-2, -1)) / math.sqrt(d_k)。

如果有掩码（mask），则将掩码矩阵中值为 0 的位置对应的注意力分数替换为 -1e9。

对注意力分数应用 Softmax 函数：attention_scores.softmax(dim=-1)。

如果有 Dropout，则对注意力分数应用 Dropout。

返回注意力加权的值和注意力分数。

前向传播：通过线性层生成查询（Q）、键（K）和值（V）。

将 Q、K、V 分割成多个头：query.view(...).transpose(1, 2)。

调用 attention 方法计算多头注意力的输出。

将多头注意力的输出拼接并线性变换：x.transpose(1, 2).contiguous().view(...); self.w_o(x)。

7. ResidualConnection 类

class ResidualConnection(nn.Module):
    def __init__(self, features: int, dropout: float):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        self.norm = LayerNormalization(features=features)    def forward(self, x, sublayer):
        # 此处和论文中 transformer 给的图略有不同, 先进行 norm 再进行 self-attention 或 feed-forward
        return x + self.dropout(sublayer(self.norm(x)))

功能：实现残差连接，将输入和子层的输出相加，有助于缓解梯度消失问题。

初始化：features：输入和输出的特征维度。

dropout：Dropout 概率，用于防止过拟合。

self.norm：层归一化层。

前向传播：对输入应用层归一化：self.norm(x)。

将归一化后的输入传递给子层：sublayer(self.norm(x))。

对 子层的输出应用 Dropout：self.dropout(...)。

将输入和子层的输出相加：x + ...。

8. EncoderBlock 类

class EncoderBlock(nn.Module):
    def __init__(self, features: int, self_attention_block: MultiHeadAttentionBlock,
                 feed_forward_block: FeedForwardBlock, dropout: float):
        super().__init__()
        self.self_attention_block = self_attention_block
        self.feed_forward_block = feed_forward_block
        self.residual_connections = nn.ModuleList(
            [ResidualConnection(features, dropout) for _ in range(2)]
        )    def forward(self, x, src_mask):
        x = self.residual_connections[0](x, lambda x: self.self_attention_block(x, x, x, src_mask))
        x = self.residual_connections[1](x, self.feed_forward_block)
        return x

功能：实现编码器块，包含一个自注意力层和一个前馈神经网络层，以及两个残差连接。

初始化：features：输入和输出的特征维度。

self_attention_block：多头自注意力块。

feed_forward_block：前馈神经网络块。

self.residual_connections：包含两个残差连接的模块列表。

前向传播：对输入应用自注意力层和第一个残差连接：self.residual_connections[0](x, lambda x: self.self_attention_block(x, x, x, src_mask))。

对自注意力层的输出应用前馈神经网络层和第二个残差连接：self.residual_connections[1](x, self.feed_forward_block)。

9. Encoder 类

class Encoder(nn.Module):
    def __init__(self, features: int, layers: nn.ModuleList):
        super().__init__()
        self.layers = layers
        self.norm = LayerNormalization(features=features)    def forward(self, x, mask):
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

功能：实现编码器，由多个编码器块组成，最后进行层归一化。

初始化：features：输入和输出的特征维度。

layers：包含多个编码器块的模块列表。

self.norm：层归一化层。

前向传播：依次将输入传递给每个编码器块：for layer in self.layers: x = layer(x, mask)。

对最后一个编码器块的输出应用层归一化：self.norm(x)。

10. DecoderBlock 类

class DecoderBlock(nn.Module):
    def __init__(self, features: int, self_attention_block: MultiHeadAttentionBlock,
                 cross_attention_block: MultiHeadAttentionBlock,
                 feed_forward_block: FeedForwardBlock, dropout: float):
        super().__init__()
        self.self_attention_block = self_attention_block
        self.cross_attention_block = cross_attention_block
        self.feed_forward_block = feed_forward_block
        self.residual_connections = nn.ModuleList(
            [ResidualConnection(features, dropout) for _ in range(3)]
        )    def forward(self, x, encoder_output, src_mask, tgt_mask):
        x = self.residual_connections[0](x, lambda x: self.self_attention_block(x, x, x, tgt_mask))
        x = self.residual_connections[1](x, lambda x: self.cross_attention_block(x, encoder_output, encoder_output,
                                                                                 src_mask))
        x = self.residual_connections[2](x, self.feed_forward_block)
        return x

功能：实现解码器块，包含一个自注意力层、一个交叉注意力层和一个前馈神经网络层，以及三个残差连接。

初始化：features：输入和输出的特征维度。

self_attention_block：多头自注意力块。

cross_attention_block：多头交叉注意力块。

feed_forward_block：前馈神经网络块。

self.residual_connections：包含三个残差连接的模块列表。

前向传播：对输入应用自注意力层和第一个残差连接：self.residual_connections[0](x, lambda x: self.self_attention_block(x, x, x, tgt_mask))，这里的 tgt_mask 用于屏蔽目标序列中后续位置的信息，避免解码器提前看到未来的信息。

对自注意力层的输出应用交叉注意力层和第二个残差连接：self.residual_connections[1](x, lambda x: self.cross_attention_block(x, encoder_output, encoder_output, src_mask))，交叉注意力层中，查询来自解码器当前层的输入，键和值来自编码器的输出，src_mask 用于屏蔽源序列中填充位置的信息。

对交叉注意力层的输出应用前馈神经网络层和第三个残差连接：self.residual_connections[2](x, self.feed_forward_block)。

11. Decoder 类

class Decoder(nn.Module):
    def __init__(self, features: int, layers: nn.ModuleList):
        super().__init__()
        self.layers = layers
        self.norm = LayerNormalization(features=features)    def forward(self, x, encoder_output, src_mask, tgt_mask):
        for layer in self.layers:
            x = layer(x, encoder_output, src_mask, tgt_mask)
        return self.norm(x)

功能：实现解码器，由多个解码器块组成，最后进行层归一化。

初始化：features：输入和输出的特征维度。

layers：包含多个解码器块的模块列表。

self.norm：层归一化层。

前向传播：依次将输入、编码器输出、源序列掩码和目标序列掩码传递给每个解码器块：for layer in self.layers: x = layer(x, encoder_output, src_mask, tgt_mask)。

对最后一个解码器块的输出应用层归一化：self.norm(x)。

12. ProjectionLayer 类

class ProjectionLayer(nn.Module):
    def __init__(self, d_model, vocab_size):
        super().__init__()
        self.proj = nn.Linear(d_model, vocab_size)    def forward(self, x):
        # (batch, seq_len, d_model) --> (batch, seq_len, vocab_size)
        return self.proj(x)

功能：将解码器的输出投影到词汇表大小的维度，以便进行词的预测。

初始化：d_model：输入的特征维度。

vocab_size：目标词汇表的大小。

self.proj：线性层，将 d_model 维的输入映射到 vocab_size 维的输出。

前向传播：对输入应用线性层，将其维度从 (batch, seq_len, d_model) 转换为 (batch, seq_len, vocab_size)。

13. Transformer 类

class Transformer(nn.Module):
    def __init__(self, encoder: Encoder, decoder: Decoder, src_embed: InputEmbeddings, tgt_embed: InputEmbeddings,
                 src_pos: PositionalEncoding, tgt_pos: PositionalEncoding, projection_layer: ProjectionLayer):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = src_embed
        self.tgt_embed = tgt_embed
        self.src_pos = src_pos
        self.tgt_pos = tgt_pos
        self.projection_layer = projection_layer    def encode(self, src, src_mask):
        # (batch, seq_len, d_model)
        src = self.src_embed(src)
        src = self.src_pos(src)
        return self.encoder(src, src_mask)    def decode(self, encoder_output: torch.Tensor, src_mask: torch.Tensor, tgt: torch.Tensor, tgt_mask: torch.Tensor):
        # (batch, seq_len, d_model)
        tgt = self.tgt_embed(tgt)
        tgt = self.tgt_pos(tgt)
        return self.decoder(tgt, encoder_output, src_mask, tgt_mask)    def project(self, x):
        return self.projection_layer(x)  # (batch, seq_len, vocab_size)

功能：整合编码器、解码器、嵌入层、位置编码层和投影层，构建完整的 Transformer 模型。

初始化：encoder：编码器模块。

decoder：解码器模块。

src_embed：源语言嵌入层。

tgt_embed：目标语言嵌入层。

src_pos：源语言位置编码层。

tgt_pos：目标语言位置编码层。

projection_layer：投影层。

encode 方法：对源序列进行嵌入和位置编码：src = self.src_embed(src); src = self.src_pos(src)。

将处理后的源序列传递给编码器：self.encoder(src, src_mask)。

decode 方法：对目标序列进行嵌入和位置编码：tgt = self.tgt_embed(tgt); tgt = self.tgt_pos(tgt)。

将处理后的目标序列、编码器输出、源序列掩码和目标序列掩码传递给解码器：self.decoder(tgt, encoder_output, src_mask, tgt_mask)。

project 方法：将解码器的输出传递给投影层，得到最终的预测结果。

14. build_transformer 函数

def build_transformer(src_vocab_size: int, tgt_vocab_size: int, src_seq_len: int, tgt_seq_len: int, d_model: int = 512,
                      N: int = 6, h: int = 8, dropout: float = 0.1, d_ff: int = 2048) -> Transformer:# embedding layers
    src_embed = InputEmbeddings(d_model, src_vocab_size)
    tgt_embed = InputEmbeddings(d_model, tgt_vocab_size)# positional encoding layers
    src_pos = PositionalEncoding(d_model, src_seq_len, dropout)
    tgt_pos = PositionalEncoding(d_model, tgt_seq_len, dropout)# encoder blocks
    encoder_blocks = []for _ in range(N):
        encoder_self_attention_block = MultiHeadAttentionBlock(d_model, h, dropout)
        feed_forward_block = FeedForwardBlock(d_model, d_ff, dropout)
        encoder_block = EncoderBlock(d_model, encoder_self_attention_block, feed_forward_block, dropout)
        encoder_blocks.append(encoder_block)# encoder
    encoder = Encoder(d_model, nn.ModuleList(encoder_blocks))# decoder blocks
    decoder_blocks = []for _ in range(N):
        decoder_self_attention_block = MultiHeadAttentionBlock(d_model, h, dropout)
        decoder_cross_attention_block = MultiHeadAttentionBlock(d_model, h, dropout)
        feed_forward_block = FeedForwardBlock(d_model, d_ff, dropout)
        decoder_block = DecoderBlock(d_model, decoder_self_attention_block, decoder_cross_attention_block,
                                     feed_forward_block, dropout)
        decoder_blocks.append(decoder_block)# decoder
    decoder = Decoder(d_model, nn.ModuleList(decoder_blocks))# projection layer
    projection_layer = ProjectionLayer(d_model, tgt_vocab_size)# transformer
    transformer = Transformer(encoder, decoder, src_embed, tgt_embed, src_pos, tgt_pos, projection_layer)# 初始化transformer模型中的参数for p in transformer.parameters():if p.dim() > 1:
            nn.init.xavier_normal_(p)return transformer

功能：根据给定的参数构建一个完整的 Transformer 模型。

参数：src_vocab_size：源语言词汇表的大小。

tgt_vocab_size：目标语言词汇表的大小。

src_seq_len：源序列的最大长度。

tgt_seq_len：目标序列的最大长度。

d_model：词向量和隐藏层的维度，默认值为 512。

N：编码器和解码器中块的数量，默认值为 6。

h：多头注意力机制中头的数量，默认值为 8。

dropout：Dropout 概率，默认值为 0.1。

d_ff：前馈神经网络中间层的维度，默认值为 2048。

构建过程：分别创建源语言和目标语言的嵌入层和位置编码层。

构建编码器块和编码器。

构建解码器块和解码器。

创建投影层。

将上述组件组合成一个完整的 Transformer 模型。

对模型的参数进行初始化，对于维度大于 1 的参数，使用 xavier_normal_ 方法进行初始化。

返回构建好的 Transformer 模型。

完整代码实现transfomer

import torch
import torch.nn as nn
import mathclass InputEmbeddings(nn.Module):# 词嵌入本质会训练一个矩阵, (vocab_size, d_model)def __init__(self, d_model: int, vocab_size: int):super().__init__()
        self.d_model = d_model
        self.vocab_size = vocab_size
        self.embedding = nn.Embedding(vocab_size, d_model)def forward(self, x):# nn.Embedding 里面是有参数的, 使用的时候相当于 lookup table# 乘上权重来自于论文 3.4return self.embedding(x) * math.sqrt(self.d_model)class PositionalEncoding(nn.Module):# d_model 是每个位置的向量维度, seq_len 是总共多少个位置, dropout 防止过拟合def __init__(self, d_model: int, seq_len: int, dropout: float):super().__init__()
        self.d_model = d_model
        self.seq_len = seq_len
        self.dropout = nn.Dropout(dropout)# 创建一个矩阵, 形状(seq_len, d_model)
        pe = torch.zeros(seq_len, d_model)# 创建一个向量 (seq_len, 1)
        position = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1)# 实现公式 3.5, 但是在 log space 下计算的
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)        pe = pe.unsqueeze(0)  # (1, seq_len, d_model)# 不需要去训练的参数, 最终它也是会保存在模型文件中的, 可以把它注册到 buffer 中
        self.register_buffer('pe', pe)def forward(self, x):# (batch, seq_len, d_model)
        x = x + (self.pe[:, :x.shape[1], :]).requires_grad_(False)return self.dropout(x)class LayerNormalization(nn.Module):# layer norm 是对每个样本不同位置的向量求均值和方差,然后再进行归一化# features 是一个样本的特征数量def __init__(self, features: int, eps: float = 1e-6):super().__init__()
        self.eps = eps
        self.alpha = nn.Parameter(torch.ones(features))
        self.bias = nn.Parameter(torch.zeros(features))def forward(self, x):# x: (batch, seq_len, hidden_size) 每条文本的每个token由多个特征构成
        mean = x.mean(dim=-1, keepdim=True)
        std = x.std(dim=-1, keepdim=True)return self.alpha * (x - mean) / (std + self.eps) + self.biasclass FeedForwardBlock(nn.Module):def __init__(self, d_model: int, d_ff: int, dropout: float):super().__init__()
        self.linear_1 = nn.Linear(d_model, d_ff)
        self.dropout = nn.Dropout(p=dropout)
        self.linear_2 = nn.Linear(d_ff, d_model)def forward(self, x):# 实现的就是论文中 3.3 FFN(x)return self.linear_2(self.dropout(torch.relu(self.linear_1(x))))class MultiHeadAttentionBlock(nn.Module):def __init__(self, d_model: int, h: int, dropout: float):super().__init__()        self.d_model = d_model
        self.h = hassert d_model % h == 0, "d_model 必须可以整除 h"# 每个 head 对应的维度是多大
        self.d_k = d_model // h        self.w_q = nn.Linear(d_model, d_model, bias=False)
        self.w_k = nn.Linear(d_model, d_model, bias=False)
        self.w_v = nn.Linear(d_model, d_model, bias=False)
        self.w_o = nn.Linear(d_model, d_model, bias=False)        self.dropout = nn.Dropout(dropout)@staticmethoddef attention(query, key, value, mask, dropout: nn.Dropout):
        d_k = query.shape[-1]        attention_scores = (query @ key.transpose(-2, -1)) / math.sqrt(d_k)if mask is not None:# 将 mask 矩阵中值为 0的位置,对应在 attention_scores 里面的值用 -1e9 替换
            attention_scores.masked_fill_(mask == 0, -1e9)
        attention_scores = attention_scores.softmax(dim=-1)  # (Batch, h, seq_len, seq_len)if dropout is not None:
            attention_scores = dropout(attention_scores)return (attention_scores @ value), attention_scoresdef forward(self, q, k, v, mask):
        query = self.w_q(q)  # Q
        key = self.w_k(k)  # K
        value = self.w_v(v)  # V# 多头注意力机制 切割 Q K V 为了给不同的 heads# (Batch, Seq_Len, d_model) --> (Batch, Seq_Len, h, d_k) --> (Batch, h, Seq_Len, d_k)
        query = query.view(query.shape[0], query.shape[1], self.h, self.d_k).transpose(1, 2)
        key = key.view(key.shape[0], key.shape[1], self.h, self.d_k).transpose(1, 2)
        value = value.view(value.shape[0], value.shape[1], self.h, self.d_k).transpose(1, 2)        x, self.attention_scores = MultiHeadAttentionBlock.attention(query, key, value, mask, self.dropout)# (Batch, h, seq_len, d_k) --> (Batch, Seq_Len, h, d_k) --> (Batch, Seq_Len, d_model)
        x = x.transpose(1, 2).contiguous().view(x.shape[0], -1, self.h * self.d_k)return self.w_o(x)class ResidualConnection(nn.Module):def __init__(self, features: int, dropout: float):super().__init__()
        self.dropout = nn.Dropout(dropout)
        self.norm = LayerNormalization(features=features)def forward(self, x, sublayer):# 此处和论文中 transformer 给的图略有不同, 先进行 norm 再进行 self-attention 或 feed-forwardreturn x + self.dropout(sublayer(self.norm(x)))class EncoderBlock(nn.Module):def __init__(self, features: int, self_attention_block: MultiHeadAttentionBlock,
                 feed_forward_block: FeedForwardBlock, dropout: float):super().__init__()
        self.self_attention_block = self_attention_block
        self.feed_forward_block = feed_forward_block
        self.residual_connections = nn.ModuleList([ResidualConnection(features, dropout) for _ in range(2)])def forward(self, x, src_mask):
        x = self.residual_connections[0](x, lambda x: self.self_attention_block(x, x, x, src_mask))
        x = self.residual_connections[1](x, self.feed_forward_block)return xclass Encoder(nn.Module):def __init__(self, features: int, layers: nn.ModuleList):super().__init__()
        self.layers = layers
        self.norm = LayerNormalization(features=features)def forward(self, x, mask):for layer in self.layers:
            x = layer(x, mask)return self.norm(x)class DecoderBlock(nn.Module):def __init__(self, features: int, self_attention_block: MultiHeadAttentionBlock,
                 cross_attention_block: MultiHeadAttentionBlock,
                 feed_forward_block: FeedForwardBlock, dropout: float):super().__init__()
        self.self_attention_block = self_attention_block
        self.cross_attention_block = cross_attention_block
        self.feed_forward_block = feed_forward_block
        self.residual_connections = nn.ModuleList([ResidualConnection(features, dropout) for _ in range(3)])def forward(self, x, encoder_output, src_mask, tgt_mask):
        x = self.residual_connections[0](x, lambda x: self.self_attention_block(x, x, x, tgt_mask))
        x = self.residual_connections[1](x, lambda x: self.cross_attention_block(x, encoder_output, encoder_output,
                                                                                 src_mask))
        x = self.residual_connections[2](x, self.feed_forward_block)return xclass Decoder(nn.Module):def __init__(self, features: int, layers: nn.ModuleList):super().__init__()
        self.layers = layers
        self.norm = LayerNormalization(features=features)def forward(self, x, encoder_output, src_mask, tgt_mask):for layer in self.layers:
            x = layer(x, encoder_output, src_mask, tgt_mask)return self.norm(x)class ProjectionLayer(nn.Module):def __init__(self, d_model, vocab_size):super().__init__()
        self.proj = nn.Linear(d_model, vocab_size)def forward(self, x):# (batch, seq_len, d_model) --> (batch, seq_len, vocab_size)return self.proj(x)class Transformer(nn.Module):def __init__(self, encoder: Encoder, decoder: Decoder, src_embed: InputEmbeddings, tgt_embed: InputEmbeddings,
                 src_pos: PositionalEncoding, tgt_pos: PositionalEncoding, projection_layer: ProjectionLayer):super().__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = src_embed
        self.tgt_embed = tgt_embed
        self.src_pos = src_pos
        self.tgt_pos = tgt_pos
        self.projection_layer = projection_layerdef encode(self, src, src_mask):# (batch, seq_len, d_model)
        src = self.src_embed(src)
        src = self.src_pos(src)return self.encoder(src, src_mask)def decode(self, encoder_output: torch.Tensor, src_mask: torch.Tensor, tgt: torch.Tensor, tgt_mask: torch.Tensor):# (batch, seq_len, d_model)
        tgt = self.tgt_embed(tgt)
        tgt = self.tgt_pos(tgt)return self.decoder(tgt, encoder_output, src_mask, tgt_mask)def project(self, x):return self.projection_layer(x)  # (batch, seq_len, vocab_size)def build_transformer(src_vocab_size: int, tgt_vocab_size: int, src_seq_len: int, tgt_seq_len: int, d_model: int = 512,
                      N: int = 6, h: int = 8, dropout: float = 0.1, d_ff: int = 2048) -> Transformer:# embedding layers
    src_embed = InputEmbeddings(d_model, src_vocab_size)
    tgt_embed = InputEmbeddings(d_model, tgt_vocab_size)# positional encoding layers
    src_pos = PositionalEncoding(d_model, src_seq_len, dropout)
    tgt_pos = PositionalEncoding(d_model, tgt_seq_len, dropout)# encoder blocks
    encoder_blocks = []for _ in range(N):
        encoder_self_attention_block = MultiHeadAttentionBlock(d_model, h, dropout)
        feed_forward_block = FeedForwardBlock(d_model, d_ff, dropout)
        encoder_block = EncoderBlock(d_model, encoder_self_attention_block, feed_forward_block, dropout)
        encoder_blocks.append(encoder_block)# encoder
    encoder = Encoder(d_model, nn.ModuleList(encoder_blocks))# decoder blocks
    decoder_blocks = []for _ in range(N):
        decoder_self_attention_block = MultiHeadAttentionBlock(d_model, h, dropout)
        decoder_cross_attention_block = MultiHeadAttentionBlock(d_model, h, dropout)
        feed_forward_block = FeedForwardBlock(d_model, d_ff, dropout)
        decoder_block = DecoderBlock(d_model, decoder_self_attention_block, decoder_cross_attention_block,
                                     feed_forward_block, dropout)
        decoder_blocks.append(decoder_block)# decoder
    decoder = Decoder(d_model, nn.ModuleList(decoder_blocks))# projection layer
    projection_layer = ProjectionLayer(d_model, tgt_vocab_size)# transformer
    transformer = Transformer(encoder, decoder, src_embed, tgt_embed, src_pos, tgt_pos, projection_layer)# 初始化transformer模型中的参数for p in transformer.parameters():if p.dim() > 1:
            nn.init.xavier_normal_(p)return transformer