Transformer细节（六）——详解Transformer各层结构和组成及代码实现

多学学多写写

已于 2024-07-05 12:37:49 修改

阅读量1.4k

点赞数 49

文章标签： transformer 深度学习人工智能

于 2024-06-23 09:55:51 首次发布

本文链接：https://blog.csdn.net/weixin_47129891/article/details/139883769

版权

Transformer 模型的架构是由多个编码器（Encoder）和解码器（Decoder）层堆叠而成的。

一、编码器（Encoder）

编码器由多个相同的编码器层（Encoder Layer）堆叠而成。每个编码器层包含两个主要子层：自注意力（Self-Attention）子层和前馈神经网络（Feed Forward Neural Network，FFN）子层。

1. 自注意力子层（Self-Attention Layer）

输入与输出维度

输入张量的形状为 \((N, L, d_{\text{model}})\)

计算过程

注意力机制：通过查询（Query）、键（Key）和值（Value）来计算注意力得分，公式为：

\[ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V \]

多头注意力（Multi-Head Attention）

将输入分成多头，每个头独立计算注意力，再将结果拼接并通过线性变换。多头注意力公式为：\[ \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)W^O \]
其中，每个头的计算为：\[ \text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V) \]

层归一化（Layer Normalization）

在自注意力操作之后进行归一化，并添加残差连接（Residual Connection）。

2. 前馈神经网络子层（Feed Forward Neural Network, FFN）

结构

由两个线性变换和一个激活函数（通常是ReLU）组成，公式为：\[ \text{FFN}(x) = \max(0, xW_1 + b_1)W_2 + b_2 \]

层归一化（Layer Normalization）

在前馈神经网络之后进行归一化，并添加残差连接。

3. 输出维度

每个编码器层的输出维度仍然是 \((N, L, d_{\text{model}})\)。

二、解码器（Decoder）

解码器也由多个相同的解码器层（Decoder Layer）堆叠而成。解码器层在编码器层的基础上增加了一个编码器-解码器注意力（Encoder-Decoder Attention）子层。

1. 掩码自注意力子层（Masked Self-Attention Layer）

类似于编码器的自注意力子层，但在计算注意力得分时使用了掩码（mask），以防止当前位置看到未来的信息。

2. 编码器-解码器注意力子层（Encoder-Decoder Attention Layer）

该子层的输入来自编码器层的输出和解码器层的前一层输出。计算过程与自注意力相同，但查询（Query）来自解码器，键（Key）和值（Value）来自编码器。

3. 前馈神经网络子层（Feed Forward Neural Network, FFN）

与编码器层中的前馈神经网络相同。

4. 输出维度

每个解码器层的输出维度也是 \((N, L, d_{\text{model}})\)。

三、Transformer 模型的整体结构

1. 输入嵌入（Input Embedding）

输入序列首先通过嵌入层（Embedding Layer）映射为 \((N, L, d_{\text{model}})\) 的张量。
添加位置编码（Positional Encoding），以便模型捕获序列中的位置信息。

2. 编码器堆叠（Stack of Encoders）

输入经过多个编码器层的处理，得到编码后的序列表示。

3. 解码器堆叠（Stack of Decoders）

解码器接收编码器的输出和目标序列的嵌入，经过多个解码器层的处理，生成最终的输出序列。

4. 输出层（Output Layer）

解码器的输出通过线性变换和 softmax 层，生成每个位置的预测分布。

四、代码实现

import torch
import torch.nn as nn
import torch.nn.functional as F

class Transformer(nn.Module):
    def __init__(self, num_layers, d_model, num_heads, d_ff):
        super(Transformer, self).__init__()
        self.encoder_layers = nn.ModuleList([EncoderLayer(d_model, num_heads, d_ff) for _ in range(num_layers)])
        self.decoder_layers = nn.ModuleList([DecoderLayer(d_model, num_heads, d_ff) for _ in range(num_layers)])
        
    def forward(self, src, tgt, src_mask, tgt_mask):
        enc_output = src
        
        # 编码器层的前向传播
        for layer in self.encoder_layers:
            enc_output = layer(enc_output, src_mask)
        
        dec_output = tgt
        
        # 解码器层的前向传播
        for layer in self.decoder_layers:
            dec_output = layer(dec_output, enc_output, src_mask, tgt_mask)
        
        return dec_output

# 简化的编码器层和解码器层
class EncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff):
        super(EncoderLayer, self).__init__()
        self.self_attn = nn.MultiheadAttention(d_model, num_heads)
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model)
        )
        self.layer_norm1 = nn.LayerNorm(d_model)
        self.layer_norm2 = nn.LayerNorm(d_model)

    def forward(self, x, src_mask):
        # 自注意力
        self_attn_output, _ = self.self_attn(x, x, x, attn_mask=src_mask)
        x = x + self_attn_output
        x = self.layer_norm1(x)
        
        # 前馈神经网络
        ff_output = self.feed_forward(x)
        x = x + ff_output
        x = self.layer_norm2(x)
        
        return x

class DecoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff):
        super(DecoderLayer, self).__init__()
        self.self_attn = nn.MultiheadAttention(d_model, num_heads)
        self.cross_attn = nn.MultiheadAttention(d_model, num_heads)
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model)
        )
        self.layer_norm1 = nn.LayerNorm(d_model)
        self.layer_norm2 = nn.LayerNorm(d_model)
        self.layer_norm3 = nn.LayerNorm(d_model)

    def forward(self, x, enc_output, src_mask, tgt_mask):
        # 自注意力
        self_attn_output, _ = self.self_attn(x, x, x, attn_mask=tgt_mask)
        x = x + self_attn_output
        x = self.layer_norm1(x)
        
        # 编码器-解码器注意力
        cross_attn_output, _ = self.cross_attn(x, enc_output, enc_output, attn_mask=src_mask)
        x = x + cross_attn_output
        x = self.layer_norm2(x)
        
        # 前馈神经网络
        ff_output = self.feed_forward(x)
        x = x + ff_output
        x = self.layer_norm3(x)
        
        return x

# 示例参数
num_layers = 6
d_model = 512
num_heads = 8
d_ff = 2048
sequence_length = 10
batch_size = 32

# 初始化Transformer模型
transformer = Transformer(num_layers, d_model, num_heads, d_ff)

# 生成示例输入
src = torch.rand(sequence_length, batch_size, d_model)  # 编码器输入
tgt = torch.rand(sequence_length, batch_size, d_model)  # 解码器输入
src_mask = torch.ones(sequence_length, sequence_length).bool()  # 编码器掩码
tgt_mask = torch.ones(sequence_length, sequence_length).bool()  # 解码器掩码

# 前向传播
output = transformer(src, tgt, src_mask, tgt_mask)

print(output.shape)  # 应输出: torch.Size([10, 32, 512])

五、总结

这种架构设计使得 Transformer 模型可以高效并行地处理序列数据，并且能够捕获序列中的长距离依赖关系。具体如何并行处理数据，参考我之前的博客Transformer细节（三）——Transformer是并行or顺序处理数据？-CSDN博客