使用Pytorch从零开始构建Transformer

在本教程中，我们将使用 PyTorch 从头开始构建一个基本的 Transformer 模型。Vaswani 等人提出的 Transformer 模型。在论文“Attention is All You Need”中，是一种专为序列到序列任务（例如机器翻译和文本摘要）而设计的深度学习架构。它基于自注意力机制，已成为许多最先进的自然语言处理模型（如 GPT 和 BERT）的基础。

要详细了解 Transformer 模型，请访问这两篇文章：

All you need to know about ‘Attention’ and ‘Transformers’ — In-depth Understanding — Part 1
All you need to know about ‘Attention’ and ‘Transformers’ — In-depth Understanding — Part 2

要构建 Transformer 模型，我们将按照以下步骤操作：

导入必要的库和模块；
定义基本构建块：多头注意力、位置前馈网络、位置编码；
构建编码器和解码器层；
组合编码器和解码器层以创建完整的Transformer 模型；
准备样本数据；
训练模型

让我们首先导入必要的库和模块。

import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as data
import math
import copy

现在，我们将定义 Transformer 模型的基本构建块。

多头注意力

在这里插入图片描述
多头注意力机制计算序列中每对位置之间的注意力。它由多个“注意力头”组成，捕获输入序列的不同方面。

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
        
    def scaled_dot_product_attention(self, Q, K, V, mask=None):
        attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            attn_scores = attn_scores.masked_fill(mask == 0, -1e9)
        attn_probs = torch.softmax(attn_scores, dim=-1)
        output = torch.matmul(attn_probs, V)
        return output
        
    def split_heads(self, x):
        batch_size, seq_length, d_model = x.size()
        return x.view(batch_size, seq_length, self.num_heads, self.d_k).transpose(1, 2)
        
    def combine_heads(self, x):
        batch_size, _, seq_length, d_k = x.size()
        return x.transpose(1, 2).contiguous().view(batch_size, seq_length, self.d_model)
        
    def forward(self, Q, K, V, mask=None):
        Q = self.split_heads(self.W_q(Q))
        K = self.split_heads(self.W_k(K))
        V = self.split_heads(self.W_v(V))
        
        attn_output = self.scaled_dot_product_attention(Q, K, V, mask)
        output = self.W_o(self.combine_heads(attn_output))
        return output

MultiHeadAttention 代码使用输入参数和线性变换层初始化模块。它计算注意力分数，将输入张量重塑为多个头，并组合所有头的注意力输出。前向方法计算多头自注意力，使模型能够关注输入序列的某些不同方面。

位置式前馈网络

class PositionWiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff):
        super(PositionWiseFeedForward, self).__init__()
        self.fc1 = nn.Linear(d_model, d_ff)
        self.fc2 = nn.Linear(d_ff, d_model)
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

PositionWiseFeedForward 类扩展了 PyTorch 的 nn.Module 并实现了位置明智的前馈网络。该类使用两个线性变换层和一个 ReLU 激活函数进行初始化。前向方法依次应用这些变换和激活函数来计算输出。此过程使模型能够在进行预测时考虑输入元素的位置。

位置编码

位置编码用于注入输入序列中每个标记的位置信息。它使用不同频率的正弦和余弦函数来生成位置编码。

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_seq_length):
        super(PositionalEncoding, self).__init__()
        
        pe = torch.zeros(max_seq_length, d_model)
        position = torch.arange(0, max_seq_length, dtype=torch.float).unsqueeze(1)
        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)
        
        self.register_buffer('pe', pe.unsqueeze(0))
        
    def forward(self, x):
        return x + self.pe[:, :x.size(1)]

PositionalEncoding 类使用输入参数 d_model 和 max_seq_length 进行初始化，创建一个张量来存储位置编码值。该类根据缩放因子 div_term 分别计算偶数和奇数索引的正弦和余弦值。前向方法通过将存储的位置编码值添加到输入张量来计算位置编码，从而使模型能够捕获输入序列的位置信息。

现在，我们将构建编码器和解码器层。

编码器层

在这里插入图片描述
编码器层由多头注意力层、位置前馈层和两个层归一化层组成。

class EncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout):
        super(EncoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.feed_forward = PositionWiseFeedForward(d_model, d_ff)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x, mask):
        attn_output = self.self_attn(x, x, x, mask)
        x = self.norm1(x + self.dropout(attn_output))
        ff_output = self.feed_forward(x)
        x = self.norm2(x + self.dropout(ff_output))
        return x

EncoderLayer 类使用输入参数和组件进行初始化，包括 MultiHeadAttention 模块、PositionWiseFeedForward 模块、两层归一化模块和 dropout 层。前向方法通过应用自注意力、将注意力输出添加到输入张量并对结果进行归一化来计算编码器层输出。然后，它计算位置前馈输出，将其与归一化的自注意力输出相结合，并在返回处理后的张量之前对最终结果进行归一化。

解码层

在这里插入图片描述
解码器层由两个多头注意力层、一个位置前馈层和三个层归一化层组成。

class DecoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout):
        super(DecoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.cross_attn = MultiHeadAttention(d_model, num_heads)
        self.feed_forward = PositionWiseFeedForward(d_model, d_ff)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x, enc_output, src_mask, tgt_mask):
        attn_output = self.self_attn(x, x, x, tgt_mask)
        x = self.norm1(x + self.dropout(attn_output))
        attn_output = self.cross_attn(x, enc_output, enc_output, src_mask)
        x = self.norm2(x + self.dropout(attn_output))
        ff_output = self.feed_forward(x)
        x = self.norm3(x + self.dropout(ff_output))
        return x

DecoderLayer 使用输入参数和组件进行初始化，例如用于屏蔽自注意力和交叉注意力的 MultiHeadAttention 模块、PositionWiseFeedForward 模块、三层归一化模块和 dropout 层。

前向方法通过执行以下步骤来计算解码器层输出：

计算屏蔽自注意力输出并将其添加到输入张量，然后进行 dropout 和层归一化。
计算解码器和编码器输出之间的交叉注意力输出，并将其添加到归一化屏蔽自注意力输出中，然后进行 dropout 和层归一化。
计算位置前馈输出并将其与归一化交叉注意力输出相结合，然后进行 dropout 和层归一化。
返回处理后的张量。
这些操作使解码器能够基于输入和编码器输出生成目标序列。

现在，让我们组合编码器层和解码器层来创建完整的 Transformer 模型。

Transformer Model

在这里插入图片描述

将它们合并在一起：

class Transformer(nn.Module):
    def __init__(self, src_vocab_size, tgt_vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_length, dropout):
        super(Transformer, self).__init__()
        self.encoder_embedding = nn.Embedding(src_vocab_size, d_model)
        self.decoder_embedding = nn.Embedding(tgt_vocab_size, d_model)
        self.positional_encoding = PositionalEncoding(d_model, max_seq_length)

        self.encoder_layers = nn.ModuleList([EncoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)])
        self.decoder_layers = nn.ModuleList([DecoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)])

        self.fc = nn.Linear(d_model, tgt_vocab_size)
        self.dropout = nn.Dropout(dropout)

    def generate_mask(self, src, tgt):
        src_mask = (src != 0).unsqueeze(1).unsqueeze(2)
        tgt_mask = (tgt != 0).unsqueeze(1).unsqueeze(3)
        seq_length = tgt.size(1)
        nopeak_mask = (1 - torch.triu(torch.ones(1, seq_length, seq_length), diagonal=1)).bool()
        tgt_mask = tgt_mask & nopeak_mask
        return src_mask, tgt_mask

    def forward(self, src, tgt):
        src_mask, tgt_mask = self.generate_mask(src, tgt)
        src_embedded = self.dropout(self.positional_encoding(self.encoder_embedding(src)))
        tgt_embedded = self.dropout(self.positional_encoding(self.decoder_embedding(tgt)))

        enc_output = src_embedded
        for enc_layer in self.encoder_layers:
            enc_output = enc_layer(enc_output, src_mask)

        dec_output = tgt_embedded
        for dec_layer in self.decoder_layers:
            dec_output = dec_layer(dec_output, enc_output, src_mask, tgt_mask)

        output = self.fc(dec_output)
        return output

Transformer 类组合了前面定义的模块来创建完整的 Transformer 模型。在初始化期间，Transformer 模块设置输入参数并初始化各种组件，包括源序列和目标序列的嵌入层、PositionalEncoding 模块、用于创建堆叠层的 EncoderLayer 和 DecoderLayer 模块、用于投影解码器输出的线性层以及 dropout 层。

generate_mask 方法为源序列和目标序列创建二进制掩码，以忽略填充标记并防止解码器关注未来的标记。forward方法通过以下步骤计算Transformer模型的输出：

使用generate_mask方法生成源和目标掩码。
计算源嵌入和目标嵌入，并应用位置编码和丢失。
通过编码器层处理源序列，更新enc_output 张量。
使用 enc_output 和掩码通过解码器层处理目标序列，并更新 dec_output 张量。
将线性投影层应用于解码器输出，获得输出 logits。

这些步骤使 Transformer 模型能够处理输入序列并根据其组件的组合功能生成输出序列。

准备样本数据

在此示例中，我们将创建一个玩具数据集用于演示目的。在实践中，您将使用更大的数据集，预处理文本，并为源语言和目标语言创建词汇映射。

src_vocab_size = 5000
tgt_vocab_size = 5000
d_model = 512
num_heads = 8
num_layers = 6
d_ff = 2048
max_seq_length = 100
dropout = 0.1

transformer = Transformer(src_vocab_size, tgt_vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_length, dropout)

# Generate random sample data
src_data = torch.randint(1, src_vocab_size, (64, max_seq_length))  # (batch_size, seq_length)
tgt_data = torch.randint(1, tgt_vocab_size, (64, max_seq_length))  # (batch_size, seq_length)

训练模型

现在我们将使用样本数据训练模型。在实践中，您将使用更大的数据集并将其拆分为训练集和验证集。

criterion = nn.CrossEntropyLoss(ignore_index=0)
optimizer = optim.Adam(transformer.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9)

transformer.train()

for epoch in range(100):
    optimizer.zero_grad()
    output = transformer(src_data, tgt_data[:, :-1])
    loss = criterion(output.contiguous().view(-1, tgt_vocab_size), tgt_data[:, 1:].contiguous().view(-1))
    loss.backward()
    optimizer.step()
    print(f"Epoch: {epoch+1}, Loss: {loss.item()}")

我们可以使用这种方式在 Pytorch 中从头开始构建一个简单的 Transformer。所有大型语言模型都使用这些 Transformer 编码器或解码器块进行训练。因此，了解基础Transformer网络非常重要。希望这篇文章对所有想要深入研究大语言模型(LLM)的人有所帮助。

参考文献

Attention is all you need
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. Gomez, {. Kaiser, and I. Polosukhin. Advances in Neural Information Processing Systems , page 5998–6008. (2017)