transformer 代码详解

Transformer是一种基于自注意力机制的神经网络模型，最初被提出用于自然语言处理任务，如机器翻译和文本生成。它的核心思想是利用自注意力机制来捕捉输入序列中不同位置之间的依赖关系，从而更好地理解输入序列。

Transformer模型由编码器和解码器两部分组成，其中编码器用于将输入序列转换为一系列隐藏表示，解码器则用于根据这些隐藏表示生成输出序列。编码器和解码器都由多个相同的层组成，每个层都包含一个多头自注意力子层和一个全连接前馈子层。

在多头自注意力子层中，输入序列被映射到一个高维空间中，并计算出每个位置与其他位置之间的相似度。然后，根据这些相似度计算出每个位置对其他位置的权重，从而得到一个加权平均值作为该位置的表示。这种自注意力机制可以捕捉到输入序列中不同位置之间的依赖关系，从而更好地理解输入序列。

在全连接前馈子层中，每个位置的表示被传递到一个全连接神经网络中进行处理，以进一步提取特征。

下面是一个简单的Transformer模型的代码实现：

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

class Transformer(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_layers, num_heads):
        super(Transformer, self).__init__()
        self.embedding = nn.Embedding(input_dim, hidden_dim)
        self.encoder_layers = nn.ModuleList([
            nn.TransformerEncoderLayer(hidden_dim, num_heads) for _ in range(num_layers)
        ])
        self.decoder_layers = nn.ModuleList([
            nn.TransformerDecoderLayer(hidden_dim, num_heads) for _ in range(num_layers)
        ])
        self.fc = nn.Linear(hidden_dim, input_dim)

    def forward(self, src, tgt):
        src_embedded = self.embedding(src)
        tgt_embedded = self.embedding(tgt)
        src_mask = self.generate_square_subsequent_mask(src.shape[1])
        tgt_mask = self.generate_square_subsequent_mask(tgt.shape[1])
        memory = self.encode(src_embedded, src_mask)
        output = self.decode(tgt_embedded, memory, tgt_mask)
        output = self.fc(output)
        return output

    def encode(self, src_embedded, src_mask):
        output = src_embedded
        for layer in self.encoder_layers:
            output = layer(output, src_mask)
        return output

    def decode(self, tgt_embedded, memory, tgt_mask):
        output = tgt_embedded
        for layer in self.decoder_layers:
            output = layer(output, memory, tgt_mask)
        return output

    def generate_square_subsequent_mask(self, size):
        mask = (torch.triu(torch.ones(size, size)) == 1).transpose(0, 1)
        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
        return mask

这个代码实现了一个简单的Transformer模型，包括编码器和解码器。其中，`input_dim`表示输入序列的维度，`hidden_dim`表示隐藏层的维度，`num_layers`表示编码器和解码器的层数，`num_heads`表示多头注意力机制中的头数。在`forward`方法中，首先将输入序列和目标序列分别进行嵌入，然后生成对应的掩码，接着将输入序列传递给编码器进行编码，得到一个表示整个输入序列的隐藏表示，最后将目标序列和隐藏表示传递给解码器进行解码，得到最终的输出序列。最后，通过一个全连接层将输出序列映射到指定维度的空间中。