Transformer简单编码实现

不懂就要问

已于 2024-11-23 23:07:19 修改

阅读量1.2k

点赞数 34

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

于 2024-11-23 21:44:52 首次发布

本文链接：https://blog.youkuaiyun.com/weixin_54171657/article/details/143998777

版权

一、前言：

Transformer 架构是一种深度学习模型，最初由 Vaswani 等人在 2017 年的论文《Attention Is All You Need》中提出，主要用于处理序列到序列（sequence-to-sequence）的任务，如机器翻译、文本摘要等。它的核心思想是使用注意力机制（Attention Mechanism）来解决传统序列处理模型（如循环神经网络 RNN）在处理长距离依赖问题时的局限性。

Transformer 架构的一些特点：

1：编码器-解码器结构：

编码器（Encoder）：由多个相同的层（Layer）组成，每个层包含两个主要的子层：自注意力（Self-Attention）机制和前馈神经网络（Feed-Forward Neural Network）。编码器的主要任务是理解输入序列。
解码器（Decoder）：同样由多个相同的层组成，每个层包含三个主要的子层：自注意力机制、编码器-解码器注意力（Encoder-Decoder Attention）机制和前馈神经网络。解码器的主要任务是生成输出序列。

2：自注意力机制（Self-Attention）：

允许模型在序列的不同位置之间直接捕捉依赖关系，无论这些位置之间的距离有多远。
通过计算序列中每个元素对其他元素的注意力权重，来确定每个元素在序列中的上下文关系。

3：并行处理能力：

由于自注意力机制不依赖于序列的特定顺序，Transformer 可以并行处理整个序列，这与 RNN 的逐元素处理方式相比，大大提高了计算效率。

4：位置编码（Positional Encoding）：

由于 Transformer 不像 RNN 那样具有递归结构，因此需要一种方式来编码序列中元素的位置信息。
位置编码通常是通过添加一个与输入序列长度相关的向量来实现的，这个向量通过正弦和余弦函数生成，以保持序列中元素的顺序信息。

5：残差连接（Residual Connection）：

每个子层的输出都会加上该子层的输入，然后进行层归一化（Layer Normalization）。
这种设计有助于避免在深层网络中出现的梯度消失问题。

6：层归一化（Layer Normalization）：

在每个子层的输出上应用归一化，有助于稳定训练过程，加速收敛。

二、各部分实现：

1、位置编码：

由于自注意力机制不依赖于序列的特定顺序，Transformer 可以并行处理整个序列，这与 RNN 的逐元素处理方式相比，大大提高了计算效率。但与之对应的，对于Transformer来说：“我爱你”和“你爱我”的效果一样，但实际我们知道，这并不一样，因此，我们需要位置编码来标识个词元的位置关系，这里采用最简单的实现方式：一维绝对位置编码。

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

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).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):
        return self.dropout(x + self.pe[:, :x.size(1)].detach())

2、注意力机制

这部分是transformer架构的核心，有自注意力机制、交叉注意力机制、多头注意力机制等。

注意力机制如图：

输入三个矩阵：Q、K、V，当 Q=K=V时是自注意力机制，当Q！=K=V时，是交叉注意力机制，在trransformer架构中，编码器使用自注意力机制，解码器使用交叉注意力机制，而且都是用多个注意力头，这是为了能多层次的表示特征。

这一部分原理说明请自行探索，这里直接给出实现代码：

class MultiHeadedAttention(nn.Module):
    def __init__(self, h, d_model, dropout):
        super(MultiHeadedAttention, self).__init__()
        assert d_model % h == 0
        self.d_k = d_model // h
        self.h = h
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.attn = None
        self.dropout = nn.Dropout(dropout)

    def attention(self, q, k, v, mask=None, dropout=None):
        d_k = q.size(-1)
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        p_attn = F.softmax(scores, dim=-1)
        if dropout is not None:
            p_attn = dropout(p_attn)
        return torch.matmul(p_attn, v), p_attn

    def forward(self, query, key, value, mask=None):
        query, key, value = self.linears[0](query), self.linears[1](key), self.linears[2](value)
        if mask is not None:
            mask = mask.unsqueeze(1)
        batch_size = query.size(0)
        query = query.view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
        value = value.view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
        key = key.view(batch_size, -1, self.h, self.d_k).transpose(1, 2)

        x, self.attn = self.attention(query, key, value, mask, dropout=self.dropout)
        x = x.transpose(1, 2).view(batch_size, -1, self.h * self.d_k)
        return self.linears[-1](x)

该部分输入的Q、K、V要先经过映射，之后按照下面的公式计算即可，其中张量形状变化需要注意：输入Q、K、V的形状为[batch_size,seq_length,d_model]，需要先将d_model转化为h*d_k(h：注意力头的数量，d_k：每个头应该处理的维度)，之后调整位置为[batch_size,h,seq_length,d_k]，当完成计算得到注意力得分x和注意力矩阵sttn之后，需要将x由形状[batch_size,h,seq_length,d_k]调整为[batch_size,seq_length,d_model]，之后再经过映射返回。

这里的mask需要着重说明，在transformer架构中编码器和解码器需要用到的mask不同，在编码器中，由于同一个批次输入的序列长度都一致，这其中，有些序列是经过填充得到的，由于计算注意力机制代价比较大，那么这些填充的0就没有计算的必要，那么就要把这些填充的值用一个很小的数来替换，这样的话经过softmax就会把这些值变成0，方便后续计算。那么编码器使用的mask怎么得到的呢：

def create_encoder_mask(src_seq, pad_idx):
    """
    创建编码器的Padding Mask

    Args:
    - src_seq: 源序列 [batch_size, seq_len]
    - pad_idx: 填充标记的索引

    Returns:
    - mask: 布尔掩码 [batch_size, 1, 1, seq_len]
    """
    # 创建掩码，将填充位置设为0，非填充位置设为1
    mask = (src_seq != pad_idx).unsqueeze(1).unsqueeze(2)
    return mask

解码器主要用于生成任务，在实际情况中，生成任务是有顺序的，只有得到前一个输出才能进行下一步的计算，但是由于Transformer的注意力机制会看到全局的信息，不符合我们的要求，所以我们需要“屏蔽”掉注意力矩阵中当前位置之后的所有元素（其实就是得到一个三角矩阵），确保当前位置无法“看到”未来的信息，从而保持序列生成的顺序性，同时为了不计算那些经过填充的位置，还需要将上述矩阵与类似编码器那样的矩阵做逻辑与运算：

def create_decoder_mask(tgt_seq, pad_idx):
    """
    创建解码器的Padding Mask

    Args:
    - tgt_seq: 目标序列 [batch_size, seq_len]
    - pad_idx: 填充标记的索引

    Returns:
    - padding_mask: 填充掩码 [batch_size, 1, 1, seq_len]
    - subsequent_mask: 后续掩码 [1, seq_len, seq_len]
    """
    # 创建Padding Mask
    padding_mask = (tgt_seq != pad_idx).unsqueeze(1).unsqueeze(2)

    # 创建Subsequent Mask
    seq_len = tgt_seq.size(1)
    subsequent_mask = torch.tril(torch.ones(seq_len, seq_len)).unsqueeze(0)

    # 组合两种Mask
    decoder_mask = padding_mask & subsequent_mask.bool()

    return decoder_mask

3、前馈神经网络：

class PositionwiseFeedForward(nn.Module):
    """Implements FFN equation."""

    def __init__(self, d_model, d_ff, dropout=0.1):
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.w_2(self.dropout(F.relu(self.w_1(x))))

4、子层连接：

上面三个图是Transformer架构的三个模块，他们都有一个共同点：一个功能模块后跟残差连接，这样的话我们可以设计实现这种模块，方便复用：

class SublayerConnection(nn.Module):
    def __init__(self, d_model, dropout):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        return x + self.dropout(sublayer(self.norm(x)))

5、编码器：

编码器由N个编码器层组成，我们设计实现这个编码器层，一层包括两个子模块：自注意力模块和前馈神经网络模块：

class EncoderLayer(nn.Module):
    def __init__(self, d_model, attn, feed_forward, dropout):
        super(EncoderLayer, self).__init__()
        self.attn = attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(d_model, dropout), 2)
        self.d_model = d_model

    def forward(self, x, mask):
        # mask:标识padding的mask:[batch_size,1,1,src_seq_len]
        # 首先经过自注意力机制，再经过前馈连接网络
        x = self.sublayer[0](x, lambda x: self.attn(x, x, x, mask))
        x = self.sublayer[1](x, self.feed_forward)
        return x

N个编码器层堆叠组成编码器：

class Encoder(nn.Module):
    def __init__(self, layer, N):
        super(Encoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.d_model)

    def forward(self, x, mask):
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

其实很多模型都是只有编码器而没有解码器，例如如果你做分类、实体识别这种的而非生成任务就可以只用编码器，例如BERT和VIT。

6、解码器：

解码器由N个解码器层组成，我们设计实现这个解码器层，一层包括三个子模块：掩码自注意力模块、交叉注意力模块和前馈神经网络模块：

class DecoderLayer(nn.Module):
    def __init__(self, d_model, src_attn, self_attn, feed_forword, dropout):
        #src_attn：交叉注意力
        #self._attn：自注意力
        super(DecoderLayer, self).__init__()
        self.self_attn = self_attn
        self.feed_forward = feed_forword
        self.src_attn = src_attn
        self.sublayer = clones(SublayerConnection(d_model, dropout), 3)
        self.d_model = d_model

    def forward(self, x, m, src_mask, target_mask):
        # src_mask:处理编码器的输出，标识padding的mask：[batch_size,1,1,src_seq_len]
        # target_mask：处理解码器输入，是标识padding的mask和下三角矩阵的&:[batch_size,1,target_seq_len,target_seq_len]
        # 先经过解码器输入的自注意力机制，再经过（q=编码器输出、k和v=解码器输入）交叉注意力机制，最后经过前馈连接网络
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, target_mask))
        x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))
        x = self.sublayer[2](x, self.feed_forward)
        return x

这里需要注意两个不同的注意力模块的输入。同理，N个解码器层堆叠组成解码器：

class Decoder(nn.Module):
    def __init__(self, layer, N):
        super(Decoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.d_model)

    def forward(self, x, m, src_mask, target_mask):
        for layer in self.layers:
            x = layer(x, m, src_mask, target_mask)
        return self.norm(x)

7、生成器：

生成器很简单，就是一个映射层加一个softmax层：

class Generator(nn.Module):
    def __init__(self, d_model, vocab):
        super(Generator, self).__init__()
        self.proj = nn.Linear(d_model, vocab)

    def forward(self, x):
        
        return F.softmax(self.proj(x), dim=-1)

三、整体实现：

对于两个嵌入输入:src_embed和tgt_embed，首先src_embed经过编码器，编码器结果与tgt_embed一起传入解码器，解码器结果输入生成器得到最后的生成结果：

class transformer(nn.Module):
    def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
        super(transformer, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = src_embed
        self.tgt_embed = tgt_embed
        self.generator = generator

    def encode(self, src, src_mask):
        return self.encoder(self.src_embed(src), src_mask)

    def decode(self, m, target, src_mask, tgt_mask):
        return self.decoder(self.tgt_embed(target), m, src_mask, tgt_mask)

    def forward(self, src, target, src_mask, target_mask):
        m = self.encode(src, src_mask)
        x = self.decode(m, target, src_mask, target_mask)
        return self.generator(x)

构造模型：

这里给定参数：

src_vocab：源输入的文本词汇量大小

tgt_vocab：生成目标的文本词汇量大小（最后的输出是目标文本词汇索引）

n：编码器和解码器的层数

d_model：词向量嵌入的维度

d_ff：中间层的维度（前馈神经网络部分使用d_model*4）

h：注意力模块头的数量

dropout：用于模型泛化，丢弃参数比例

max_len：模型最长输入序列长度

def make_model(src_vocab, tgt_vocab, n=12, d_model=768, d_ff=3072, h=8, dropout=0.1, max_len=5000):
    c = copy.deepcopy
    attn = MultiHeadedAttention(h, d_model, dropout)
    ff = PositionwiseFeedForward(d_model, d_ff, dropout)
    position = PositionalEncoding(d_model, dropout, max_len)
    model = transformer(
        Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), n),
        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), n),
        nn.Sequential(Embeddings(d_model, src_vocab), c(position)),
        nn.Sequential(Embeddings(d_model, tgt_vocab), c(position)),
        Generator(d_model, tgt_vocab)
    )
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)
    return model