TR3复现Tramsformer

🍨 本文为🔗365天深度学习训练营中的学习记录博客
🍖 原作者：K同学啊

前言

Transformer模型是深度学习中的一个革命性架构，自从在NLP领域引入以来，就因其高效处理序列数据的能力而迅速成为主流。本文将通过代码实现详细剖析Transformer模型的各个组件，包括多头注意力机制、前馈神经网络、位置编码、编码器和解码器等部分。

1. Transformer模型简介

Transformer模型由Vaswani等人在2017年提出，首次在自然语言处理任务中完全摆脱了循环神经网络（RNN），依赖于自注意力机制来处理序列数据。它不仅在机器翻译、文本生成等任务中表现优异，还在各种任务中展现了良好的扩展性和性能。

Transformer的核心思想是通过自注意力机制（Self-Attention）来捕捉序列中词与词之间的关系，并通过多头注意力（Multi-Head Attention）和前馈神经网络（Feedforward Neural Network）来进一步处理这些关系。

2. Transformer的核心组件

2.1 多头注意力机制

在Transformer中，多头注意力机制（Multi-Head Attention）是最重要的组件之一。它通过对输入序列的不同部分进行多次并行的注意力计算，从而捕捉到更多的上下文信息。以下是多头注意力机制的实现：

class MultiHeadAttention(nn.Module):
    def __init__(self, hid_dim, n_heads):
        super(MultiHeadAttention, self).__init__()
        self.hid_dim = hid_dim
        self.n_heads = n_heads

        # 确保hid_dim可以被n_heads整除
        assert hid_dim % n_heads == 0

        # 定义线性变换矩阵
        self.w_q = nn.Linear(hid_dim, hid_dim)
        self.w_k = nn.Linear(hid_dim, hid_dim)
        self.w_v = nn.Linear(hid_dim, hid_dim)
        self.fc  = nn.Linear(hid_dim, hid_dim)

        # 缩放因子
        self.scale = torch.sqrt(torch.FloatTensor([hid_dim // n_heads]))

    def forward(self, query, key, value, mask=None):
        bsz = query.shape[0]
        Q = self.w_q(query)
        K = self.w_k(key)
        V = self.w_v(value)

        # 将Q, K, V拆分成多个头
        Q = Q.view(bsz, -1, self.n_heads, self.hid_dim // self.n_heads).permute(0, 2, 1, 3)
        K = K.view(bsz, -1, self.n_heads, self.hid_dim // self.n_heads).permute(0, 2, 1, 3)
        V = V.view(bsz, -1, self.n_heads, self.hid_dim // self.n_heads).permute(0, 2, 1, 3)

        # 计算注意力得分
        attention = torch.matmul(Q, K.permute(0, 1, 3, 2)) / self.scale

        if mask is not None:
            attention = attention.masked_fill(mask == 0, -1e10)

        attention = torch.softmax(attention, dim=-1)

        # 计算多头注意力的输出
        x = torch.matmul(attention, V)

        # 拼接多个头的输出
        x = x.permute(0, 2, 1, 3).contiguous()
        x = x.view(bsz, -1, self.n_heads * (self.hid_dim // self.n_heads))
        x = self.fc(x)
        return x

在这个实现中，MultiHeadAttention类首先对输入的query、key、value进行线性变换，然后将它们拆分为多个注意力头，并分别计算每个头的注意力得分。最后，将所有头的结果拼接起来并通过线性层输出。

2.2 前馈神经网络

前馈神经网络（Feedforward Neural Network）是Transformer中的另一核心组件。它通常由两层线性变换和一个ReLU激活函数组成：

class Feedforward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        super(Feedforward, self).__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(d_ff, d_model)

    def forward(self, x):
        x = torch.nn.functional.relu(self.linear1(x))
        x = self.dropout(x)
        x = self.linear2(x)
        return x

这个前馈网络将每个位置的表示独立地通过一个全连接层映射到更高维空间，再映射回原来的维度，从而增强模型的表达能力。

2.3 位置编码

由于Transformer模型不再使用RNN来处理序列数据，因此需要一种方法让模型感知到输入序列中词的顺序信息。为此，引入了位置编码（Positional Encoding）：

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model).to(device)
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) * -(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)
        self.register_buffer("pe", pe)

    def forward(self, x):
        x = x + self.pe[:, :x.size(1)].requires_grad_(False)
        return self.dropout(x)

位置编码通过正弦和余弦函数生成固定的编码，能够为不同位置的词提供唯一的表示。

3. 编码器和解码器

Transformer模型的编码器和解码器分别由多个层堆叠而成。每一层都包含一个多头注意力机制和一个前馈神经网络。

3.1 编码器层

编码器层的实现如下：

class EncoderLayer(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
        super(EncoderLayer, self).__init__()
        self.self_attn   = MultiHeadAttention(d_model, n_heads)
        self.feedforward = Feedforward(d_model, d_ff, dropout)
        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 = x + self.dropout(attn_output)
        x = self.norm1(x)

        ff_output = self.feedforward(x)
        x = x + self.dropout(ff_output)
        x = self.norm2(x)

        return x

每个编码器层首先通过多头自注意力机制处理输入序列，然后通过前馈神经网络进一步处理。最后，通过残差连接和LayerNorm层来规范化输出。

3.2 解码器层

解码器层的实现如下：

class DecoderLayer(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
        super(DecoderLayer, self).__init__()
        self.self_attn   = MultiHeadAttention(d_model, n_heads)
        self.enc_attn    = MultiHeadAttention(d_model, n_heads)
        self.feedforward = Feedforward(d_model, d_ff, dropout)
        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, self_mask, context_mask):
        attn_output = self.self_attn(x, x, x, self_mask)
        x           = x + self.dropout(attn_output)
        x           = self.norm1(x)

        attn_output = self.enc_attn(x, enc_output, enc_output, context_mask)
        x           = x + self.dropout(attn_output)
        x           = self.norm2(x)

        ff_output = self.feedforward(x)
        x = x + self.dropout(ff_output)
        x = self.norm3(x)

        return x

解码器层包含三个部分：自注意力机制、编码器-解码器注意力机制，以及前馈神经网络。解码器层不仅需要关注目标序列的自身信息，还需要从编码器的输出中提取上下文信息。

4. Transformer模型的整体架构

最后，我们将所有组件组合成完整的Transformer模型：

class Transformer(nn.Module):
    def __init__(self, vocab_size, d_model, n

_heads, n_encoder_layers, n_decoder_layers, d_ff, dropout=0.1):
        super(Transformer, self).__init__()
        self.embedding           = nn.Embedding(vocab_size, d_model)
        self.positional_encoding = PositionalEncoding(d_model, dropout)
        self.encoder_layers      = nn.ModuleList([EncoderLayer(d_model, n_heads, d_ff, dropout) for _ in range(n_encoder_layers)])
        self.decoder_layers      = nn.ModuleList([DecoderLayer(d_model, n_heads, d_ff, dropout) for _ in range(n_decoder_layers)])
        self.fc_out              = nn.Linear(d_model, vocab_size)
        self.dropout             = nn.Dropout(dropout)

    def forward(self, src, trg, src_mask, trg_mask):
        src = self.embedding(src)
        src = self.positional_encoding(src)
        trg = self.embedding(trg)
        trg = self.positional_encoding(trg)

        for layer in self.encoder_layers:
            src = layer(src, src_mask)

        for layer in self.decoder_layers:
            trg = layer(trg, src, trg_mask, src_mask)

        output = self.fc_out(trg)

        return output

在这个完整的Transformer模型中，我们首先对源语言和目标语言进行嵌入并加上位置编码。然后，经过多层编码器和解码器的处理，最后通过一个线性层输出最终的预测结果。

结果

# 使用示例
vocab_size = 10000  # 假设词汇表大小为10000
d_model    = 512
n_heads    = 8
n_encoder_layers = 6
n_decoder_layers = 6
d_ff             = 2048
dropout          = 0.1

transformer_model = Transformer(vocab_size, d_model, n_heads, n_encoder_layers, n_decoder_layers, d_ff, dropout)

# 定义输入，这里的输入是假设的，需要根据实际情况修改
src = torch.randint(0, vocab_size, (32, 10))  # 源语言句子
trg = torch.randint(0, vocab_size, (32, 20))  # 目标语言句子
src_mask = (src != 0).unsqueeze(1).unsqueeze(2)  # 掩码，用于屏蔽填充的位置
trg_mask = (trg != 0).unsqueeze(1).unsqueeze(2)  # 掩码，用于屏蔽填充的位置

# 模型前向传播
output = transformer_model(src, trg, src_mask, trg_mask)
print(output.shape)