【NLP251】Transformer API调用

1. nn.Transformer

nn.Transformer封装了Transformer中的包含编码器（Encoder）和解码器（Decoder）。如下图所示，它对Encoder和Decoder两部分的包装，它并没有实现输入中的Embedding和Positional Encoding和最后输出的Linear+Softmax部分。

我们可以注意到，在torch.nn的模块中，提供了构建transformer所需的各类元素，但却没有直接给出一个完整的Transformer算法模型。原因在于自然语言处理（NLP）领域中，不同的任务对Transformer架构有着多样化的需求。编码器和解码器作为Transformer的核心组成部分，被设计为既可以独立工作，也可以联合使用。这种灵活的设计允许研究者和开发者根据具体NLP任务的需求，将它们以不同的方式组合起来，从而适应各种复杂的应用场景。

只使用编码器的任务

编码器部分的任务是从输入数据中提取特征。编码器通常用于不需要生成新文本序列的任务。

只使用解码器的任务

解码器部分专注于生成新文本。通常使用自回归方式，基于之前的输出生成下一个词。

编码器和解码器都使用的任务：

当任务涉及到理解输入序列并生成新的输出序列时，编码器和解码器会联合使用。

这种模块化的方法不仅提高了模型的灵活性，还为研究者探索和优化特定任务的架构提供了便利。

2. nn.Transformer的参数构造详解

参数名	解释
d_model	Encoder和Decoder输入参数的特征维度。也就是词向量的维度。默认为512
nhead	多头注意力机制中，head的数量。默认值为8。
num_encoder_layers	TransformerEncoderLayer的数量。该值越大，网络越深，网络参数量越多，计算量越大。默认值为6
num_decoder_layers	TransformerDecoderLayer的数量。该值越大，网络越深，网络参数量越多，计算量越大。默认值为6
dim_feedforward	Feed Forward层（Attention后面的全连接网络）的隐藏层的神经元数量。该值越大，网络参数量越多，计算量越大。默认值为2048
dropout	内dropout值。默认值为0.1
activation	内Feed Forward层的激活函数。取值可以是string(“relu” or “gelu”)或者一个一元可调用的函数。默认值是relu
custom_encoder	自定义Encoder。若你不想用官方实现的TransformerEncoder，你可以自己实现一个。默认值为None
custom_decoder	自定义Decoder。若你不想用官方实现的TransformerDecoder，你可以自己实现一个。
layer_norm_eps	Add&Norm层中，BatchNorm的eps参数值。默认为1e-5
batch_first	batch维度是否是第一个。如果为True，则输入的shape应为(batch_size, 词数，词向量维度)，否则应为(词数, batch_size, 词向量维度)。默认为False。这个要特别注意，因为大部分人的习惯都是将batch_size放在最前面，而这个参数的默认值又是False，所以会报错。
norm_first	是否要先执行norm。例如，在图中的执行顺序为 Attention -> Add -> Norm。若该值为True，则执行顺序变为：Norm -> Attention -> Add。

3. 构建Transformer网络时常用API

我们在1.小节中提到nn.Transformer封装了Transformer中的包含编码器（Encoder）和解码器（Decoder）它并没有实现输入中的Embedding和Positional Encoding和最后输出的Linear+Softmax部分，所以在使用TransformerAPI时我们常使用以下API

3.1分割的编码器与解码器

nn.TransformerEncoderLayer和nn.TransformerDecoderLayer：这两个类分别代表Transformer架构中编码器和解码器的单层结构。它们内部包含了自注意力机制（self-attention）和前馈神经网络（feedforward network），同时还具备必要的层归一化（layer normalization）和残差连接（residual connections），以确保模型在训练过程中的稳定性和性能表现。

nn.TransformerEncoder与nn.TransformerDecoder: 这两个类是Transformer编码器的实现和解码器的实现，其中nn.TransformerEncoder包含了多个nn.TransformerEncoderLayer层的堆叠，nn.TransformerDecoder包含了多个nn.TransformerDecoderLayer层的堆叠。

3.2 其他API组件

nn.MultiheadAttention: 这个模块实现了多头注意力机制，这是Transformer模型的核心组件之一。多头注意力允许模型在不同的位置同时处理来自序列不同部分的信息，这有助于捕捉序列内的复杂依赖关系。

nn.LayerNorm: 层归一化（Layer Normalization）通常用在Transformer的各个子层的输出上，有助于稳定训练过程，并且提高了训练的速度和效果。

nn.PositionalEncoding: 虽然PyTorch的nn模块中没有这个类，但在使用Transformer时通常会添加一个位置编码层来给模型提供关于词汇在序列中位置的信息，因为Transformer本身不具备处理序列顺序的能力。

nn.Embedding：一个预训练好的语义空间，它将每个标记（如单词、字符等）映射到一个高维空间的向量。这使得模型能够处理文本数据，并为每个唯一的标记捕获丰富的语义属性。嵌入层通常是自然语言处理模型的第一层，用于将离散的文本数据转化为连续的向量表示。其输入是索引列表，输出是对应的嵌入向量。

nn.Transformer.generate_square_subsequent_mask：掩码函数。用于生成一个方形矩阵，用作Transformer模型中自注意力机制的上三角遮罩。这个遮罩确保在序列生成任务中，例如语言模型中，任何给定的元素只会考虑到序列中先于它的元素（即它只能看到过去的信息，不能看到未来的信息）。这种掩码通常在解码器部分使用，防止在预测下一个输出时“作弊”。具体来说，该函数创建了一个方阵，其中对角线及其以下的元素为0（表示可以“看到”这些位置的元素），其余元素为负无穷大（在softmax之前应用，表示位置被屏蔽，不应该有注意力权重）。

import torch.nn as nn
nn.Transformer.generate_square_subsequent_mask(5) # 5指的是target的维度

4. 文本分类任务实战（只使用编码器的任务）

import torch
from torch import nn
from torch.nn import TransformerEncoder, TransformerEncoderLayer

# 假设超参数和数据预处理已经完成
vocab_size = 1000  # 词汇量
embedding_size = 512  # 嵌入维度
nhead = 8  # 多头注意力的头数
num_encoder_layers = 6  # 编码器层数
dropout = 0.1  # dropout比率
num_classes = 10  # 类别数

class TextClassificationModel(nn.Module):
    def __init__(self, vocab_size, embedding_size, nhead, num_encoder_layers, dropout, num_classes):
        super(TextClassificationModel, self).__init__()
        #只使用了Transformer的Encoder部分
        encoder_layers = TransformerEncoderLayer(d_model=embedding_size, nhead=nhead, dropout=dropout)
        self.transformer_encoder = TransformerEncoder(encoder_layer=encoder_layers, num_layers=num_encoder_layers)

        self.embedding = nn.Embedding(vocab_size, embedding_size)
        self.positional_encoding = nn.Parameter(torch.zeros(1, 5000, embedding_size))  # 假定最大句子长度为 5000
        self.linear = nn.Linear(embedding_size, num_classes)
        
    #----必填参数----
    #src：Encoder的输入。也就是将token进行Embedding和Positional Encoding之后的tensor。必填参数。Shape为(batch size，sequence length，embedding size)
    
    #----可选参数----
    #src_mask：对src进行mask。不常用。Shape为(sequence length, sequence length)
    #src_key_padding_mask：对src的token进行mask. 常用。Shape为(batch size, sequence length) 

    def forward(self, src, src_mask, src_padding_mask):
        src_emb = self.embedding(src) + self.positional_encoding[:, :src.size(1)]
        encoded_src = self.transformer_encoder(src_emb, src_mask, src_padding_mask)
        # 假设我们只关心输入序列的第一个元素
        return self.linear(encoded_src[:, 0])

# 假定 src 已经是预处理并被转换成 tensor 的索引
# src = [batch_size, src_len]

model = TextClassificationModel(vocab_size, embedding_size, nhead, num_encoder_layers, dropout, num_classes)

# 接下来你需要定义mask和padding mask，并且运行模型
# 比如 src_mask, src_padding_mask
# 输出会是一个[batch_size, num_classes]的tensor

5. 类似GPT的文本生成式模型（只使用解码器的任务）

import torch
from torch import nn

class GPTLikeModel(nn.Module):
    def __init__(self, vocab_size, embedding_size, nhead, num_decoder_layers, dropout):
        super(GPTLikeModel, self).__init__()
        self.embedding_size = embedding_size
        self.embeddings = nn.Embedding(vocab_size, embedding_size)
        self.positional_encodings = nn.Parameter(torch.zeros(1, 1024, embedding_size))
        
        #只使用了Transformer的Decoder部分
        self.transformer_decoder_layer = nn.TransformerDecoderLayer(d_model=embedding_size,  nhead=nhead, dropout=dropout)
        self.transformer_decoder = nn.TransformerDecoder(self.transformer_decoder_layer, num_layers=num_decoder_layers)

        self.output_layer = nn.Linear(embedding_size, vocab_size)
    
    #生成上三角矩阵掩码
    def generate_square_subsequent_mask(self, sz):
        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
        return mask
    
    #----必填参数---- 
    #tgt:与src同理，Decoder的输入。 必填参数。也就是将token进行Embedding和Positional Encoding之后的tensor。Shape为(batch size，sequence length，embedding size)
    #----可选参数----
    #tgt_mask：对tgt进行mask。常用。Shape为(sequence length, sequence length) 
    #tgt_key_padding_mask：对tgt的token进行mask。常用。Shape为(batch size, sequence length) 
    #memory_key_padding_mask：对tgt的token进行mask。不常用。Shape为(batch size, sequence length) 
    def forward(self, tgt, memory_key_padding_mask):
        input_mask = self.generate_square_subsequent_mask(input.size(0)).to(input.device)
        embeddings = self.embeddings(input) * math.sqrt(self.embedding_size)
        embeddings += self.positional_encodings[:, :tgt.size(0), :]
        transformer_output = self.transformer_decoder(embeddings, memory, tgt_mask=tgt_mask)
        output = self.output_layer(transformer_output)
        return output

    def generate(self, input, max_length):
        memory = None
        generated = input
        for _ in range(max_length):
            
            #进行一个字的生成
            output = self.forward(generated, memory)
            
            #只取最后一个时间步的输出，包含所有可能生成的字的概率
            next_token_logits = output[-1, :, :]  
            
            #使用argmax取出概率最大的字
            next_token = torch.argmax(next_token_logits, dim=-1).unsqueeze(0) 
            
            #合并生成的字与之前的文本
            generated = torch.cat((generated, next_token), dim=0)
            if next_token.item() == eos_token_id:  # 假设eos_token_id是结束标记的ID
                break
        return generated

6. 机器翻译任务（编码器和解码器都使用的任务）

import torch
from torch import nn
from torch.nn import Transformer

# 假设一些超参数和数据预处理已经完成
src_vocab_size = 1000  # 源语言词汇量
tgt_vocab_size = 1000  # 目标语言词汇量
embedding_size = 512  # 嵌入维度
nhead = 8  # 多头注意力中的头数
num_encoder_layers = 6  # 编码器层数
num_decoder_layers = 6  # 解码器层数
dropout = 0.1  # dropout比率

class TranslationModel(nn.Module):
    def __init__(self, src_vocab_size, tgt_vocab_size, embedding_size, nhead, num_encoder_layers, num_decoder_layers, dropout):
        super(TranslationModel, self).__init__()
        self.transformer = Transformer(d_model=embedding_size, nhead=nhead, num_encoder_layers=num_encoder_layers, num_decoder_layers=num_decoder_layers, dropout=dropout)
        self.src_tok_emb = nn.Embedding(src_vocab_size, embedding_size)
        self.tgt_tok_emb = nn.Embedding(tgt_vocab_size, embedding_size)
        self.positional_encoding = nn.Parameter(torch.zeros(1, 5000, embedding_size))  # 假定最大句子长度为 5000
        self.generator = nn.Linear(embedding_size, tgt_vocab_size)
        
    #----必填参数----
    #src：Encoder的输入。也就是将token进行Embedding和Positional Encoding之后的tensor。必填参数。Shape为(batch size，sequence length，embedding size)
    #tgt:与src同理，Decoder的输入。 必填参数。Shape为(sequence length，embedding size)
    
    #----可选参数----
    #src_mask：对src进行mask。不常用。Shape为(sequence length, sequence length)
    #tgt_mask：对tgt进行mask。常用。Shape为(sequence length, sequence length) 
    #src_key_padding_mask：对src的token进行mask. 常用。Shape为(batch size, sequence length) 
    #tgt_key_padding_mask：对tgt的token进行mask。常用。Shape为(batch size, sequence length) 
    
    def forward(self, src, tgt, src_mask = None, tgt_mask = None, memory_mask = None, src_key_padding_mask = None, tgt_key_padding_mask = None,memory_key_padding_mask = None):
        src_emb = self.src_tok_emb(src) + self.positional_encoding[:, :src.size(1)]
        tgt_emb = self.tgt_tok_emb(tgt) + self.positional_encoding[:, :tgt.size(1)]
        outs = self.transformer(src_emb, tgt_emb, src_mask, tgt_mask, None, src_key_padding_mask, tgt_key_padding_mask, memory_key_padding_mask)
        return self.generator(outs)

# 假定 src 和 tgt 已经是预处理并被转换成 tensor 的索引
# src = [batch_size, src_len]
# tgt = [batch_size, tgt_len]

model = TranslationModel(src_vocab_size, tgt_vocab_size, embedding_size, nhead, num_encoder_layers, num_decoder_layers, dropout)

# 接下来你需要定义mask和padding mask，并且运行模型
# 比如 src_mask, tgt_mask, src_padding_mask, tgt_padding_mask, memory_key_padding_mask
# 输出会是一个[batch_size, tgt_len, tgt_vocab_size]的tensor