【PyTorch][chapter31][transformer-4]

摘要：

    我们常看到Transformer 的代码和损失函数跟论文的架构有些差异，

这里面重点讲解为什么会产生这种差异

参考

    https://www.youtube.com/watch?v=1IKrHh2X0F0

   PyTorch中实现Transformer模型 - zh-jp - 博客园

目录：

1.    Self-Attention Using Scaled Dot-Product Approach
2.    Intro to scaled dot-product attention
3.     多头注意力机制
4.    Self-Attention VS Cross-Attention
5.    PostLN,
6.    PreLN 
7.    ResiDual Transformers
8.   参考代码

一  Self-Attention Using Scaled Dot-Product Approach

          众所周知，由于近期像ChatGPT这样的改进，大型语言模型（LLMs）最近获得了极大的关注，而attention机制是这些模型的核心。目标是通过视觉表示来解释注意力机制背后的概念，以便在本系列结束时，你们能够很好地理解注意力机制和Transformer

1.1： what is attention mechanism 

      让我们从一个例子开始，我们有一个句子：“Jane is going to the cinema to watch a new _____”，后面跟着一个空格，那么空格中可以填入什么呢？

一些可能的答案

“movies”（电影）、

“comedy”（喜剧），

“a new action movie”（一部新动作片）

“new romantic movie”（一部新浪漫片）

      如果你有更多的上下文信息，比如你知道Jane喜欢喜剧，那么我们就可以很容易地预测出Jane要去看一部新喜剧。但当然，我们不能说是“book”（书）、“cat”（猫）、“grocery store”（杂货店）或“tennis”（网球）。

     这个预测完全基于上下文（context）或序列中提供的信息。从高层次上讲，这就是语言模型的训练方式和工作原理。因此，在给定前面单词的上下文的情况下，语言模型需要预测下一个单词。为了做到这一点，语言模型需要理解序列中单词之间的关系，然后基于这种关系，它可以学习在预测空格时应关注哪些单词。

     现在，让我们在图中看看这些关系，我在每对单词之间用这些箭头来表示它们

句子中的每个单词都要以不同程度的关注度去关注其他单词，换句话说，这些单词之间的相似性是基于上下文的。例如，单词“Cinema”和“watch”之间有更强的关系，这种关系可以帮助我们预测空格中的单词是“movies”。注意力模型的目标是捕捉序列中单词之间的关系，这些关系代表了上下文相似性。我们可以使用一个矩阵来表示上下文相似性，其中每一行和每一列都对应句子中的单个单词。在这个热图中，较高的上下文相似性用红色表示，较低的相似性用蓝色表示。尽管这只是一个用来解释概念的例子，并非来自实际的计算，但这里的目的是解释注意力矩阵作为句子中单词之间相似性矩阵的概念。

二  Intro to scaled dot-product attention

        目标是构建一个上下文相似性矩阵

        现在我们要研究一个模型，叫做缩放点积注意力（Scaled Dot-Product Attention），这个模型是在一篇名为《Attention is All You Need》的论文中提出的。

         注意力机制并不是这篇论文首创的，因为在它之前就已经存在，但这篇论文通过添加一个缩放因子提高了其性能。尽管这个缩放只是一个小的改进，但这篇论文因其提出的名为Transformer的特殊架构而闻名。这是第一篇在前馈神经网络中使用注意力机制的论文，因为之前的论文都是在循环神经网络（如LSTM或GRU）中使用注意力。我们不想进一步分散注意力，所以让我们来看看缩放点积注意力是如何工作的。

它基本上接收三个矩阵：一个查询矩阵Q，一个键矩阵K，和一个值矩阵V。请注意，掩码步骤是可选的，因此为了简化，我们将省略它。

主要计算流程如下：

1.1  step1 Matmul 

step2:  Scale 注意力缩放

 step3   Mask  

一、处理输入序列长度不一致的问题

在自然语言处理（NLP）任务中，输入的句子长度通常是不一致的。然而，Transformer模型需要固定的输入维度，因此较短的句子通常通过填充（padding）的方式达到与最长句子相同的长度。然而，这些填充位置的信息对于模型来说是无效的，不应该参与计算。这时，Padding Mask（填充掩码）的作用就显得尤为重要。

Padding Mask的作用是将填充的位置标记为特定的值（如0），使得模型在计算注意力时能够忽略这些填充位置的信息。通过这种方式，模型可以高效地处理不同长度的输入数据，而无需担心填充位置对模型性能的影响。

二、模拟真实场景中的信息流

在Transformer的解码过程中，模型需要按照自回归的方式进行训练，即模型在预测下一个词时，只能依赖于当前词及之前的词的信息，而不能看到未来的词。为了实现这一点，我们需要使用Sequence Mask（序列掩码）来限制自注意力机制的计算范围。

Sequence Mask通过构建一个掩码矩阵来实现，该矩阵是一个下三角矩阵，对角线及以下的部分为1（不遮挡），对角线以上的部分为0（遮挡）。这样，在计算自注意力时，模型就只能使用当前词及之前的词的信息了。通过这种方式，我们可以模拟真实场景中的信息流，使模型的预测更加符合实际。

综上所述，Transformer中的Mask操作在处理不同长度输入数据、控制自注意力机制的计算范围以及模拟真实场景中的信息流方面发挥着重要作用。这些Mask操作共同提升了Transformer模型的性能和泛化能力。

 下图对应的是Sequence Mask

step4  Softmax 

step5  输出：

       矩阵Q、K和V实际上是序列，其中每一行对应于输入序列中的一个标记。根据通用约定，我们用小写粗体符号来表示Q、K和V的行，以与矩阵本身相区分。因此，小写符号q是大小为dk的查询向量，小写符号k是键向量，其大小必须与q相同，因为我们将对q和k执行点积运算。小写符号v是值向量，其大小可以与q和k不同，因此我们用dv来表示其大小。由于这些是行向量，我们使用下标符号来表示它们的维度，例如，Q_1（表示第一个查询向量）的维度为Dk。

3  softmax

4 输出 

二   多头注意力机制

 作用：

          相对于Scaled Dot Product Attention ,

          多头注意力机制可以从输入序列不同位置的不同子空间中提取信息

三 Self-Attention VS Cross-Attention

   1:  self-attention

        编码器，解码器都采用了这种架构

       只关注 input sequence 内部词之间的关系

2：  Cross-Attention

       解码器中采用了该架构

   有两个 sequnce  , X(英语）, Y（德语）

    Y attention 计算和X 之间的relation shape

四   PostLN

    Transformer 论文中使用的是PostLN,先做残差连接 ，再LayerNorm,所以其优化器  要使用

warm-up 技术

    4.1 Post LN 架构

 4.2  PostLN 对应的优化器  Warnup

为什么需要warm-up ,下面这篇论文讨论了其原因

4.3 PostLN 问题

阻碍梯度在反向传播过程中的顺畅流动

五  PreLN

     解决了梯度消失或者梯度爆炸问题，不需要用warm-up优化器.

  先做归一化，再做残差。

5.2 优缺点

问题：

Representation Collapse（表示崩塌）是深度学习和自监督学习中的一个关键问题，特别是在对比学习中更为常见。以下是对Representation Collapse的详细解释：

一、定义

Representation Collapse指的是模型学习到的特征表示变得过于简单或退化，失去了对输入数据的有效区分能力。在极端情况下，所有输入都被映射到同一个点或非常相似的几个点上，导致模型无法区分不同的输入数据。

二、发生场景及原因

1. 对比学习中的表示崩塌：

   • 在对比学习中，如果正样本对之间的相似度被过度优化，而忽视了整体分布，就可能导致表示崩塌。
   • 如果查询编码器和键编码器完全相同，模型可能会找到一种“捷径”：将所有输入映射到一个固定的表示，从而轻易地区分正负样本对，但失去了对输入的有效编码。

2. 预训练模型Fine-tuning中的表示崩塌：

   • 对预训练模型进行Fine-tuning时，如果任务差异较大、数据不足或学习率过于激进，可能导致模型失去原始预训练模型所具有的多样性和丰富性的特征表示。
   • 任务差异大：预训练模型和目标任务领域相差较大时，Fine-tuning过程中可能导致模型丢失原有知识而无法适应新任务。
   • 数据不足：Fine-tuning阶段的数据量较小或数据分布与预训练数据差异较大时，模型可能过度依赖于少量的新数据。
   • 学习率过于激进：Fine-tuning过程中使用过大的学习率可能导致模型权重更新过于激烈，从而破坏原有的表示结构。

三、解决方法

为了避免Representation Collapse，可以采取以下方法：

1. 渐进的Fine-tuning：逐渐调整学习率，确保模型在Fine-tuning过程中平稳地适应新任务，防止权重更新过于激进。
2. 合理选择预训练模型：选择与目标任务相近的预训练模型，以减小任务之间的差异。
3. 数据增强：利用数据增强技术增加Fine-tuning阶段的数据多样性，有助于模型更好地适应新任务。
4. 特征蒸馏：使用特征蒸馏等技术，引导模型保留预训练时学到的有用特征。

四、具体技术解决方案

在自监督学习中，针对Representation Collapse问题，一些具体的技术解决方案包括：

1. BYOL（Bootstrap Your Own Latent）：通过EMA（指数移动平均）、predictor和BN（批量归一化）等方式在隐性地进行uniformity（均匀性）约束，以防止模型塌陷。
2. 数据增强和正则化：通过数据增强技术增加输入数据的多样性，并使用正则化方法防止模型过拟合。
3. 损失函数设计：设计合理的损失函数，如VICReg等，以平衡alignment（对齐）和uniformity（均匀性）之间的关系，从而避免表示崩塌。

综上所述，Representation Collapse是深度学习和自监督学习中的一个重要问题。为了解决这个问题，需要合理选择预训练模型、采用渐进的Fine-tuning策略、进行数据增强和正则化、以及设计合理的损失函数等。

六  Residual Transformer

 七  代码

  这个代码使用了PreLN 方案

# -*- coding: utf-8 -*-
"""
Created on Mon Sep  9 15:29:47 2024

@author: chengxf2
"""
import torch  
import matplotlib.pyplot as plt# -*- coding: utf-8 -*-
import torch.nn as nn  
import math
import  torch.nn.functional as F
import copy 

def draw_subsequent(mask):
    plt.figure(figsize=(5,5))
    plt.imshow(mask[0][0])
    
def get_padding_mask(seq):
    """
    编码器掩码器
    ----------
        key : [batch_size, seq_length]
          编码器中使用
    Returns
    -------
        mask : [batch_size,1,1,seq_length]
        DESCRIPTION.

    """
    mask = seq.eq(0)
    #[batch_size, 1, 1, seq_key_length]
    mask = mask.unsqueeze(1).unsqueeze(2)
    return mask

def get_subsequence_mask(t):
    """
    解码器掩码器
    Parameters
    ----------
        t : 时刻
        t+1 时刻以后全部屏蔽.

    Returns
    -------
        mask : [batch_size, nheads, t, t]
        上三角矩阵，屏蔽对应为True的地方

    """
    #掩码张量函数：上三角矩阵,主对角线上部分都是1
    #t+1 时刻都设置为true,代表需要遮掩
    mask = torch.triu(torch.ones((t,t)),1).bool()
    mask = mask.unsqueeze(0).unsqueeze(1)
    return mask


class  Embedding(nn.Module):
    #词嵌入模块
    def __init__(self, d_model=512, vocab_size=1000):
        #d_model: 词向量的维度
        #vocab_size: 词汇大小，中文翻译英文中，对应中文词汇大小和英文词汇带大小
        super(Embedding,self).__init__()
        self.layer_emb = nn.Embedding(vocab_size, d_model)
        self.d_model = d_model
    
    def forward(self, inputs):
        """  
        Args:  
            inputs: [batch,seq_length,vocab_size] 
        Returns:  
            outputs:  [batch,seq_length,d_model]
        """ 
        x = self.layer_emb(inputs)
        #乘以缩放因子 math.sqrt(d_model)
        outputs= x*math.sqrt(self.d_model)
        return outputs
    
class PositionalEncoding(nn.Module):
    #位置编码器
    def __init__(self, d_model=512, max_len=100, dropout=0.1):
        #d_model: 词向量维度
        #droput: 置0比率 
        #max_len: 句子最大长度
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(dropout)
        #初始化位置编码矩阵
        pe = torch.zeros(max_len, d_model)
        #绝对位置矩阵 [max_len,1]
        position = torch.arange(0, max_len).unsqueeze(1)
        #先取对数，再取指数 a= exp(log(a))
        div_term = torch.exp( torch.arange(0, d_model, 2)*(-math.log(10000)/d_model))
        pe[:,0::2]= torch.sin(position*div_term)
        pe[:,1::2]= torch.cos(position*div_term)
        #[batch_size, max_len,d_model]
        pe = pe.unsqueeze(0)
        #buffer 中的数据 （不可学习）：不会被优化器更新
        self.register_buffer("pe", pe)
    
    def forward(self, inputs):
        """
        Parameters
        ----------
            inputs : [batch_size, seq_length, d_model]

        Returns
        -------
            output: [batch_size, seq_length, d_model]

        """
        batch,seq_length,d_model = inputs.shape
        position = self.pe[:,0:seq_length,:]
        output = inputs+position
        return self.dropout(output)


class MultiheadAttention(nn.Module):
    #多头注意力机制
    def __init__(self, d_model=512, nheads=8, drop=0.1):
        super(MultiheadAttention,self).__init__()
        #投影后维度的大小
        #nheads 头的个数
        #d_model: 词的维度
        assert d_model%nheads ==0
        self.d_k = d_model//nheads
        self.nheads= nheads
        self.d_model = d_model
        self.W_query = nn.Linear(d_model, d_model)
        self.W_key = nn.Linear(d_model, d_model)
        self.W_value = nn.Linear(d_model, d_model)
        self.layer_concat = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(drop)
        
    def dot_product_attention(self, query,key,value,mask=None):
        """  
        Args:  
            query (Tensor): 输入特征 形状 (batch,nheads,seq_lengt,d_k) 
            key (Tensor):   输入特征 形状 (batch,nheads,seq_lengt,d_k) 
            value (Tensor): 输入特征形状 (batch,nheads,seq_lengt,d_k) 
            mask (Tensor): padding输入特征 形状 (batch,nheads,seq_lengt,d_k)  利用用Broadcast 原理进行计算
            mask(Tensor):  sequence输入特征形状 (1,1,seq_length_query,seq_length_query) 利用Broadcast 原理进行计算
        Returns:  
            Tensor:  形状 (batch,seq_length,d_model) 
        """ 
        batch_size,nheads,seq_length,d_k =query.shape
        #按照注意力公式，将query和Key 的转置进行矩阵乘法，然后除以缩放系数
        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)  
        if mask is not None:  
                scores = scores.masked_fill(mask, -1e9)  
        attention_weights = F.softmax(scores, dim=-1)  
        #可选项： Google没有该操作,包括李沫的教程和清华的教程有该操作
        attention_weights = self.dropout(attention_weights)  
        output = torch.matmul(attention_weights, value)  
        return output
    
    def qkv_projection(self,inputs,batch_size,seq_length,nheads,d_k):
        #投影到子空间
        x = inputs.view(batch_size, seq_length, nheads,d_k)
        #[batch_size, nheads, seq_length, d_k]
        x = x.transpose(1,2)
        return x
        
    def forward(self, query,key,value,mask=None):
        """  
        Args:  
           query (Tensor): 输入特征 形状 (batch,seq_length_q,d_model) 
           key (Tensor):   输入特征 形状   (batch,seq_length_k,d_model) 
           value (Tensor): 输入特征形状 (batch,seq_length_k,d_model) 
           mask (Tensor): padding输入特征 形状 (batch,1,1,seq_length_k) 利用用Broadcast 原理进行计算
           mask(Tensor):  sequence输入特征形状 (1,1,seq_length_query,seq_length_query) 利用Broadcast 原理进行计算
        Returns:  
           Tensor:  形状 (batch,seq_length,d_model) 
        """ 
        q = self.W_query(query)
        k = self.W_key(key)
        v = self.W_value(value)
        batch_size,seq_length, d_model = q.shape
        q = self.qkv_projection(q, batch_size,seq_length, self.nheads, self.d_k)
        k = self.qkv_projection(k, batch_size,seq_length, self.nheads, self.d_k)
        v = self.qkv_projection(v, batch_size,seq_length, self.nheads, self.d_k)
        #[batch_size, nheads, seq_length,d_k]
        x = self.dot_product_attention(q, k, v, mask)
        x = x.transpose(1,2).contiguous().view(batch_size, -1,self.d_model)
        output = self.layer_concat(x)
        return output

class FeedForward(nn.Module):
      #前馈全连接层 Feedforward fully connected layer
      #考虑注意力机制可能对复杂过程的拟合程度不够，通过增加两层网络来增加模型的能力
      #d_model: 词嵌入维度  d_ff 隐藏层维度 
      def  __init__(self, d_model=512, d_ff=2048, drop=0.1):
           super(FeedForward,self).__init__()
           self.layer_linear1 = nn.Linear(d_model, d_ff)
           self.layer_linear2 = nn.Linear(d_ff, d_model)
           self.dropout = nn.Dropout(drop)
      
      def forward(self, x):
          """  
          Args:  
             x (Tensor): 输入特征 形状 (batch,seq_length,d_model) 
          Returns:  
             Tensor:  形状 (batch,seq_length,d_model) 
          """ 
          h = F.relu(self.layer_linear1(x))
          output = self.layer_linear2(self.dropout(h))
          return output
          

class AddNorm(nn.Module):
    #残差连接，归一化
    #将自注意力机制（或前馈神经网络）的输出与输入相加（残差连接）
    #然后通过层归一化（Layer Normalization）来稳定输出。
    #这种结构有助于防止在深层网络中的梯度消失或梯度爆炸问题，并加速训练过程。
    def __init__(self, normalized_shape,dropout):
        super(AddNorm,self).__init__()
        self.norm = nn.LayerNorm(normalized_shape)
    
    def forward(self, x, sublayer_output):
        """  
        Args:  
           x (Tensor): 输入特征，用于与sublayer_output相加实现残差连接。  
           sublayer_output (Tensor): 子层（如自注意力层或前馈网络层）的输出。  
        Returns:  
           Tensor: 经过残差连接和层归一化后的输出。  
        """ 
        output = x+ sublayer_output
        output = self.norm(output)
        return output



class EncoderLayer(nn.Module):
    #编码器码器层
    def __init__(self, d_model=512, nheads=8, d_ff=2048,drop=0.1,normalized_shape=None):
        super(EncoderLayer, self).__init__()
        #d_model 词嵌入维度
        #nheads  投影子空间的个数
        #d_ff 前馈全连接的隐藏层维度
        #drop 遗忘率
        self.self_attention = MultiheadAttention(d_model, nheads,drop)
        self.addNorm_1 = AddNorm(normalized_shape,drop)
        self.feed_forward = FeedForward(d_model, d_ff, drop)
        self.addNorm_2 = AddNorm(normalized_shape,drop)
    
    def forward(self, x, mask):
        """
         
        Parameters
        ----------
            x : [batch,seq_length, d_model]
            mask : 掩码器

        Returns
        -------
            output : [batch,seq_length, d_model]

        """
        #第一个子层
        sublayer_output = self.self_attention(x,x,x,mask)
        x = self.addNorm_1(x, sublayer_output)
        #第二个子层
        sublayer_output = self.feed_forward(x)
        output = self.addNorm_2(x, sublayer_output)
        return output
    

def clone_module(module,N):
    """深拷贝意味着它会复制对象及其所有子对象（递归地），包括对象内部的嵌套对象。与浅拷贝（如 copy.copy）不同，
    深拷贝会创建一个新的对象实例，并且会复制对象中的所有属性，如果这些属性是对象类型（比如列表、字典、集合或其他自定义对象）的话，
    这些属性也会被深拷贝。使用 copy.deepcopy 的好处是，它可以让你完全独立于原始对象来操作复制后的对象，即使原始对象和其子对象被修改，
    也不会影响到深拷贝后的对象。
    """
    layers = nn.ModuleList([copy.deepcopy(module) for _ in range(N)])
    return layers

class Encoder(nn.Module):
      #编码器
      def __init__(self,layer, num_layers):
          super(Encoder,self).__init__()
          self.layers = clone_module(layer,num_layers)
          
      def forward(self, x, mask=None):
          """
          Parameters
          ----------
          x : [batch_size, seq_length, d_model]
          mask : [batch_size, 1,1,seq_length]

          Returns
          -------
          x : [batch_size, seq_length,d_model]

          """
          for layer in self.layers:
              x = layer(x,mask)
          return x
    
class DecoderLayer(nn.Module):
    #解码器层
    def __init__(self, d_model=512, nheads=8, d_ff=2048,drop=0.1, normalized_shape=512):
        super(DecoderLayer, self).__init__()
        #d_model: 词嵌入维度
        #nheads:  头个数 
        #d_ff 前馈全连接网络的隐藏层
        #drop 丢失率
        #normalized_shape： query的形状 
        #第一个子层： 自注意机制
        self.self_attention = MultiheadAttention(d_model, nheads,drop)
        self.addNorm_1 =  AddNorm(normalized_shape,drop)
        #第二个子层： cross-attention
        self.cross_attention = MultiheadAttention(d_model, nheads,drop)
        self.addNorm_2 = AddNorm(normalized_shape,drop)
        #第三个子层： 前馈全连接层
        self.feed_forward = FeedForward(d_model, d_ff, drop)
        self.addNorm_3 = AddNorm(normalized_shape,drop)
    
    def forward(self, x, encoder_inputs, source_mask, target_mask):
        """
        Parameters
        ----------
            x : 【batch,seq_length_query,d_model]
                解码器输入
            encoder_inputs : 【batch,seq_length_key,d_model]
                编码器输出，作为解码器的key,value
            source_mask : 编码器用的掩码张量
            target_mask : 解码器用的掩码张量
        Returns
        -------
            output : TYPE
            DESCRIPTION.

        """
        #第一个子层
        sublayer_output= self.self_attention(x,x,x,target_mask)
        query = self.addNorm_1(x, sublayer_output)
        #第二个子层
        sublayer_output = self.cross_attention(query,encoder_inputs,encoder_inputs,source_mask)
        x = self.addNorm_2(query, sublayer_output)
        #第三个子层
        sublayer_output= self.feed_forward(x)
        output = self.addNorm_3(x, sublayer_output)
        return output  

class Decoder(nn.Module):
    #解码器
    def __init__(self,layer, N):
        super(Decoder,self).__init__()
        
        self.layers = clone_module(layer,N)
    
    def forward(self, x, memory, source_mask, target_mask):
        """
        Parameters
        ----------
        x : 解码器输入.
        memory : 编码器输出.
        source_mask : 编码器掩码.
        target_mask : 解码器掩码

        Returns
        -------
        x : [batch_size, seq_length, d_model]

        """
        for layer in self.layers:
            x = layer(x, memory, source_mask, target_mask)
        return x
    

class  Generator(nn.Module):
    #解码器最后的输出，预测的值
    def __init__(self, d_model, vocab_size):
        #d_model: 词嵌入的维度
        #vocab: 代表词表的总大小
        super(Generator, self).__init__()
        self.layer = nn.Linear(d_model, vocab_size)
    
    def forward(self, x):
        #x: 上一层的输出张量
        output = self.layer(x)
        p_output = F.log_softmax(output,dim=-1)
        return p_output, output
    
class Transformer(nn.Module):
    #编码器-解码器结构
    def __init__(self,vocab_size =1000,max_seq_len=100, d_model=512, nheads=8, 
                 d_ff=2048, drop=0.1,num_layers=6, normalized_shape=512):
        """
        Parameters
        ----------
        encoder : 编码器对象
        decoder : 解码器对象
        source_embed : 源数据嵌入函数
        target_embed : 目标数据嵌入函数
        generator : 目标类别生成器对象
        """
        super(Transformer, self).__init__()
        
        encoder_layer = EncoderLayer(d_model, nheads, d_ff, drop, normalized_shape)
        decoder_layer = DecoderLayer(d_model, nheads, d_ff, drop, normalized_shape)
        source_embed =  Embedding(d_model,vocab_size)
        target_embed =  Embedding(d_model,vocab_size)
        self.encoder = Encoder(encoder_layer,  num_layers)
        self.decoder = Decoder(decoder_layer,  num_layers)
        self.source_embedding = nn.Sequential(source_embed,PositionalEncoding(d_model,max_seq_len,drop))
        self.target_embedding = nn.Sequential(target_embed,PositionalEncoding(d_model,max_seq_len,drop))
        self.generator = Generator(d_model, vocab_size)
        
    def transformer_encode(self, source, source_mask= None):
        #编码器函数
        x = self.source_embedding(source)
        print(x.shape)
        output = self.encoder(x, source_mask)
        return output
    
    def transformer_decode(self, memory, source_mask, target, target_mask ):
        #memory 编码器编码后的输出张量
        decoder_inputs = self.target_embedding(target)
        output = self.decoder(decoder_inputs, memory, source_mask, target_mask)
        return output
    
    
    def forward(self, source, target, source_mask, target_mask):
        """
        Parameters
        ----------
        source : 源数据
        target : 目标数据
        source_mask : 源数据掩码张量
        target_mask : 目标数据掩码张量

        """
        #调用编码器函数
        memory = self.transformer_encode(source, source_mask)
        output = self.transformer_decode(memory, source_mask, target, target_mask)
        return output
    
def getModel():
    vocab_size = 1000
    num_layers = 6
    d_model=512
    nheads=8
    d_ff=2048
    drop=0.1
    normalized_shape=512
    max_seq_len = 100
    #, nheads, d_ff, drop,num_layers, normalized_shape
    model = Transformer(vocab_size ,max_seq_len, d_model)
    source = target = torch.LongTensor([[100,2,421,508],
                                        [491,998,1,221]])
    
    source_mask = target_mask = torch.zeros(8,4,4)
    
    for p in model.parameters():
        if p.dim()>1:
            nn.init.xavier_uniform_(p)
    out = model(source, target,source_mask, target_mask)
    print(model)
    return model


    
getModel()

参考：

https://www.youtube.com/watch?v=RsuSOylfN2I&t=3s