【Transformer框架代码实现】

大佬喝可乐2023-12-23 3:42

Transformer

Transformer框架

注意力机制框架

导入必要的库

复制代码 
from torch import nn
import torch
import math
import torch.nn.functional as F

Input Embedding / Out Embedding

词嵌入层的实现

nn.Embedding(vocab_size, embed_dim)中

vocab_size: 词典长度，必须大于输入的单词总量

embed_dim：词向量维度

复制代码 
class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size, embed_dim):
        """
        :param vocab_size: 词典长度／维度／图片size(0)  必须大于输入的单词总量. batch_size=句子数量
        :param embed_dim: 词向量维度
        """
        super(TokenEmbedding, self).__init__()
        self.embed_dim = embed_dim
        self.embedding = nn.Embedding(vocab_size, embed_dim)

    def forward(self, x):
        # 对原始向量进行缩放，出自论文
        # [batch_size, vocab_size, embed_dim]
        return self.embedding(x.long()) * math.sqrt(self.embed_dim)

Positional Embedding

位置编码层的实现

与RNN不同，对于给定输入input = "I am fine"，RNN中会顺序读入每个单词，而在Transformer中则是所有单词同时读入，因此失去了"有序性"，对机器来说，只收到三个单词I，am，fine，其并不知道这三个单词之间的位置关系，为此需要为每个单词附加上额外的位置信息。比如这里我们可以简单的传（I, 1),(am, 2),(fine, 3)。这样不仅传入了句子的单词，还附加上单词的位置信息。

一种好的位置编码方案需要满足以下几条要求：

  • 能为每个时间步输出一个独一无二的编码

  • 不同长度的句子之间，任何两个时间步之间的距离应该保持一致

  • 模型应该能毫不费力地泛化到更长的句子，且值是有界的

  • 必须是确定性的

    class PositionEncoding(nn.Module):
    def init(self, d_model, max_len, dropout=0.1, device='cpu'):
    """
    :param max_len: 句子最大长度 default:5000
    :param d_model: 模型维度/embed_dim default:512 必须为偶数
    :param device:
    """
    super(PositionEncoding, self).init()
    self.dropout = nn.Dropout(dropout)
    pos_encoding = torch.zeros(max_len, d_model, device=device)
    # 2D: [max_len, d_model]
    pos = torch.arange(0, max_len, dtype=torch.float, device=device).unsqueeze(dim=1)
    # 2D: [max_len, 1]
    div_term = torch.exp(torch.arange(0, d_model, step=2, device=device).float() * (-math.log(10000) / d_model))
    # 1D: [d_model / 2]
    pos_encoding[:, 0::2] = torch.sin(pos * div_term)
    pos_encoding[:, 1::2] = torch.cos(pos * div_term)
    # 2D: [max_len, d_model]
    pos_encoding = pos_encoding.unsqueeze(dim=0)
    # print(pos_encoding.shape)
    # 3D: [1, max_len, d_model],　扩充batch_size维度 batch_size * max_len * d_model
    self.register_buffer('pos_encoding', pos_encoding)

    复制代码 
      def forward(self, x):
      # [batch_size, max_len, d_model]
      x = x + self.pos_encoding[:, :x.size(1)].requires_grad_(False)
      return self.dropout(x)

Transformer Embedding

词嵌入层＋位置编码层

复制代码 
class TransformerEmbedding(nn.Module):
    def __init__(self, vocab_size, max_len, d_model, drop, device):
        """

        :param vocab_size: 字典大小，字典中不同字符个数
        :param max_len:
        :param d_model:
        :param drop:
        :param drop_prob:
        :param device:
        """
        super(TransformerEmbedding, self).__init__()
        self.embed = TokenEmbedding(vocab_size, d_model)
        self.pos = PositionEncoding(d_model, max_len, drop, device)

    def forward(self, x):
        """
        :param x: padding后的vectors
        :return: embedding + pos_encoder 后的vectors; shape: [batch_size, max_len, d_model]
        """
        embed = self.embed(x)
        pos = self.pos(embed)
        return pos

调试例子，可以详细查看每一步运行的shape变化

复制代码 
if __name__ == '__main__':
    Input = [[1, 2, 3, 4], [5, 6, 2], [1, 7, 8, 9, 10, 11, 12, 13, 14, 15]]
    # 有三个句子［［＇你　好　世　界＇］，　［＇我　很　好＇］，　［＇你　不　是　真　正　的　快　乐＇］］
    # 句子长度不一致进行padding处理
    max_len = len(Input[2])
    print('The most length is {} sequences '.format(max_len))
    src_sentence = torch.tensor([])
    for i in range(len(Input)):
        src_sen = torch.tensor([Input[i]])
        src_sen = F.pad(src_sen, (0, max_len-len(Input[i])), 'constant', 0)
        src_sentence = torch.cat([src_sentence, src_sen], dim=0)
    # x = src_sentence.unsqueeze(dim=0)
    x = src_sentence
    print(x)
    # token_embed = TokenEmbedding(vocab_size=18, embed_dim=32)
    # x = token_embed(x)
    # print('token的维度为：', x.shape)
    # pos_embed = PositionEncoding(d_model=32, max_len=max_len)
    # x = pos_embed(x)
    # print(x.shape)
    trans = TransformerEmbedding(vocab_size=18, max_len=10, d_model=32, drop=0.1, device='cpu')
    x = trans(x)
    print('transformer的维度为：', x.shape)

ScaleDotProductAttention(self-attention)

单头注意力机制／自注意力机制实现

  • q: query 查询，
    -代表输入序列每个位置的上下文信息，可以看做单个位置向其他位置发出询问

  • k: key 键，
    -代表输入序列中每个位置的局部信息，每个位置将自己的局部信息编码为键，可以与其他位置进行比较和交互

  • v: value 值，
    -代表序列中每个位置的具体信息，可以获得具体数据
    查询会和键进行一个相似度计算，然后根据相似度的权重对值进行加权求和

    class ScaleDotProductAttention(nn.Module):
    def init(self, dim_model):
    """
    self_attention: softmax(QKT / sqrt(d_model))V
    :param dim_model: 嵌入维度，default: 512
    """
    super(ScaleDotProductAttention, self).init()
    self.q = nn.Linear(dim_model, dim_model)
    self.k = nn.Linear(dim_model, dim_model)
    self.v = nn.Linear(dim_model, dim_model)
    self.softmax = nn.Softmax(dim=-1)

    复制代码 
      def forward(self, query, key, value, mask=None):
      """
      :param query: [batch_size, max_len, d_model]
      :param key:
      :param value:
      :param mask: [max_len, max_len]
      :return:
      """
      q = self.q(query)
      k = self.k(key)
      v = self.v(value)
      d_k = query.size(-1)
      qk = self.softmax(torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(d_k))
      if mask is not None:
          qk = qk.masked_fill(mask == 0, -1e9)
      attn_weights = self.softmax(qk)
      output = torch.matmul(attn_weights, v)
      return attn_weights, output

MultiHeadAttention 多头注意力机制

多头注意力机制实现

模型在对当前位置的信息进行编码时，会过度的将注意力集中于自身的位置， 因此作者提出了通过多头注意力机制来解决这一问题。同时，使用多头注意力机制还能够给予注意力层的输出包含有不同子空间中的编码表示信息，从而增强模型的表达能力。

复制代码 
import torch
from torch import nn
import torch.nn.functional as F
import math


class MultiHeadAttention(nn.Module):
    def __init__(self, dim_model, n_head, bias=True):
        """
        :param dim_model: 词嵌入的维度，也就是前面的d_model参数，论文中的默认值为512
        :param n_head: self-attention的个数  多头注意力机制中多头的数量，也就是前面的n_head参数， 论文默认值为 8
        :param bias: 最后对多头的注意力（组合）输出进行线性变换时，是否使用偏置
        """
        super(MultiHeadAttention, self).__init__()
        self.d_model = dim_model
        self.head_dim = dim_model // n_head
        self.k_dim = self.head_dim
        self.v_dim = self.head_dim
        self.nums_head = n_head

        # self.dropout = nn.Dropout(dropout)
        assert self.nums_head * self.head_dim == self.d_model, 'embed_dim 除以 num_heads必须为整数'
        self.q_weight = nn.Parameter(torch.Tensor(dim_model, dim_model))
        # d_model = k_dim * nums_head
        self.k_weight = nn.Parameter(torch.Tensor(dim_model, dim_model))
        self.v_weight = nn.Parameter(torch.Tensor(dim_model, dim_model))
        self.output_weights = nn.Parameter(torch.Tensor(dim_model, dim_model))
        # self.output_bias = bias
        self.output = nn.Linear(dim_model, dim_model, bias=bias)

    def forward(self, query, key, value, drop_p, training=True, attn_mask=None, key_padding_mask=None):
        """
        在论文中，编码时query, key, value 都是同一个输入，
        解码时 输入的部分也都是同一个输入，
        解码和编码交互时 key,value指的是 memory, query指的是tgt
        :param drop_p: attention_weights  layer
        :param training:  是否开启训练模式
        :param query: [batch_size, tgt_len, embed_dim], tgt_len 表示目标序列的长度
        :param key:   [batch_size, src_len, embed_dim], src_len 表示源序列的长度
        :param value: [batch_size, src_len, embed_dim], src_len 表示源序列的长度
        :param attn_mask: [tgt_len,src_len] or [num_heads*batch_size,tgt_len, src_len]
                一般只在解码时使用，为了并行一次喂入所有解码部分的输入，所以要用mask来进行掩盖当前时刻之后的位置信息
        :param key_padding_mask: [batch_size, src_len], src_len 表示源序列的长度. 填充部分掩码操作
        :return:
        attn_output: [tgt_len, batch_size, embed_dim]
        attn_output_weights: # [batch_size, tgt_len, src_len]
        """
        # 第一阶段： 计算得到Q、K、V
        q = F.linear(query, self.q_weight)
        # F.linear(inputs, weights, bias)
        # query:[batch_size, src_len/tgt_len, d_model], q_weight:[d_model, d_model]
        k = F.linear(key, self.k_weight)
        v = F.linear(value, self.v_weight)
        # 第二阶段：缩放，以及attn_mask维度判断
        b_s, tgt_len, embed_dim = query.size()
        source_len = key.size(1)
        scaling = 1 / math.sqrt(self.head_dim)
        # q*scaling: [batch_size, query_len, k_dim*nums_head]
        q = q * scaling
        # print(tgt_len, source_len)

        if attn_mask is not None:
            # [tgt_len, src_len] or [nums_head*batch_size, tgt_len, src_len]
            if attn_mask.dim() == 2:
                attn_mask = attn_mask.unsqueeze(0)
                # print(attn_mask.shape)
                if list(attn_mask.size()) != [1, tgt_len, source_len]:
                    raise RuntimeError('The size of the 2D attn_mask is not correct.')
            elif attn_mask.dim() == 3:
                if list(attn_mask.size()) != [b_s*self.nums_head, tgt_len, source_len]:
                    raise RuntimeError('The size of the 3D attn_mask is not correct.')

        # 第三阶段： 计算得到注意力权重矩阵
        # [batch_size, n_head, tgt_dim, k_dim] 计算时候是［tgt_dim, k_dim］* [k_dim, tgt_dim]->
        # [batch_size, n_head, tgt_dim, tgt_dim] == [batch_size * n_head, tgt_dim, tgt_dim]
        # q = q.view(b_s, tgt_len, n_head, head_dim).transpose(1, 2)
        q = q.contiguous().view(b_s*self.nums_head, tgt_len, self.head_dim)
        # [batch_size * n_head, tgt_len, k_dim]
        k = k.contiguous().view(b_s*self.nums_head, -1, self.k_dim)
        v = v.contiguous().view(b_s*self.nums_head, -1, self.v_dim)
        # print(q.shape, k.shape, v.shape)

        attn_output_wights = torch.bmm(q, k.transpose(-1, -2))
        # [batch_size*nums_head, tgt_len, head_dim] * [batch_size*nums_head, k_dim, src_len]
        # [batch_size*nums_head, tgt_len, src_len]

        # 第四阶段： 进行相关掩码操作
        if attn_mask is not None:
            attn_output_wights = attn_output_wights + attn_mask
            # [batch_size*nums_head, tgt_len, src_len] + [batch_size*nums_head, tgt_len, src_len]
        if key_padding_mask is not None:
            attn_output_wights = attn_output_wights.view(b_s, self.nums_head, tgt_len, source_len)
            attn_output_wights = attn_output_wights.masked_fill(key_padding_mask.unsqueeze(1).unsqueeze(2) == 0, float('-inf'))
            attn_output_wights = attn_output_wights.view(b_s * self.nums_head, tgt_len, source_len)

        attn_output_wights = F.softmax(attn_output_wights, dim=-1)
        attn_output_wights = F.dropout(attn_output_wights, p=drop_p, training=training)

        attn_output = torch.bmm(attn_output_wights, v)
        # [batch_size * nums_head, tgt_len, src_len], [batch_size*nums_head, src_len, v_dim]
        # [batch_size * nums_head, tgt_len, v_dim]
        attn_output = attn_output.contiguous().view(b_s, tgt_len, self.d_model)
        # print(attn_output.shape)
        attn_output_wights = attn_output_wights.view(b_s, self.nums_head, tgt_len, source_len)

        Z = F.linear(attn_output, self.output.weight)
        return Z, attn_output, attn_output_wights.sum(dim=1) / self.nums_head

if __name__ == '__main__':
    inputs = 'transformer_embedding_out'
    src_len = 5
    batch_size = 2
    d_model = 32
    num_head = 1
    src = torch.rand(([batch_size, src_len, d_model]))
    mha = MultiHeadAttention(dim_model=d_model, n_head=num_head, bias=True)
    out, attn_out, attn_weight = mha(src, src, src, drop_p=0.1, attn_mask=None, key_padding_mask=None)
    print(out.shape, attn_out.shape, attn_weight.shape)
    # print(src, out, attn_out, attn_weight)
    # out = mha(inputs, inputs, inputs, drop_p=0.1, attn_mask=None, key_padding_mask=None)

EncoderLayer　编码层

单个编码层的实现

注意力机制经过残差连接后，进行层归一化后输出

复制代码 
class EncoderLayer(nn.Module):
    def __init__(self, embed_dim, n_head, feedforward_dim, drop_fc=None, drop_attn=None, drop_feed=None):
        """
        单个的编码层，论文中6个
        :param embed_dim:　嵌入维度，词向量和位置编码的嵌入维度，论文中：512
        :param n_head:　　　多头注意力机制个数，论文中：8个，　512//64＝8，batch_size=64
        :param drop_attn:　注意力机制层输出的随机丢弃
        :param drop_feed:　前馈网络层输出的随机丢弃
        """
        super(EncoderLayer, self).__init__()
        # TransformerEmbedding(vocab_size=2048, d_model=512, drop=0.1, drop_prob=0.1, max_len=None, device='cpu')
        # 注意力机制层　　［batch_size, max_len, d_model］
        self.attn_encoder = MultiHeadAttention(dim_model=embed_dim, n_head=n_head, bias=True)
        # 注意力机制输出进行drop
        self.drop_out1 = nn.Dropout(drop_attn)
        self.norm1 = nn.LayerNorm(embed_dim)
        # 前馈网络层
        self.feedforward = PositionWiseFeedForward(dim_model=embed_dim, drop_fc=drop_fc, dim_feedforward=feedforward_dim)
        # 前馈网络输出进行drop
        self.drop_out2 = nn.Dropout(drop_feed)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.activation = nn.ReLU()

    def forward(self, src_sentence, drop=None, training=True, attn_mask=None, src_key_padding_mask=None):
        """
        填充后的句子输入，经过词向量＋位置编码／注意力机制层／残差连接／前馈传播层／残差连接／输出
        :param src_key_padding_mask: 输入句子的填充部分做掩盖处理
        :param attn_mask: 注意力的掩码，一般在解码部分为了掩盖后续position需要使用
        :param training: 是否训练模式
        :param drop: 注意力层weights输出的drop
        :param src_sentence: 填充后的原始句子
        :return: 编码层的输出，　解码层的部分输入
        """
        # input_embedding = self.input_embedding(src_sentence)
        input_embedding = src_sentence
        attn_encoder = self.attn_encoder(input_embedding, input_embedding, input_embedding, drop_p=drop,
                                         training=training, attn_mask=attn_mask, key_padding_mask=src_key_padding_mask)[0]
        # add & norm 残差连接后归一化处理，归一化防止数值过大后续训练时候造成梯度丢失或爆炸
        embedding_attn = input_embedding + self.drop_out1(attn_encoder)
        attn_out = self.norm1(embedding_attn)
        attn_out = self.activation(attn_out)
        # 前馈连接，Feed_Forward层
        feed_forward = self.feedforward(attn_out)
        # Feed_Forward层后的残差连接＋归一化
        attn_feed = attn_encoder + self.drop_out2(feed_forward)
        encoder_layer_out = self.norm2(attn_feed)

        return encoder_layer_out

Encoder多层编码块／前馈网络层

多个编码层的实现

复制代码 
class Encoder(nn.Module):
    def __init__(self, embed_dim, n_head, feedforward_dim, drop_fc, drop_attn, drop_feed, n_layers):
        super(Encoder, self).__init__()
        self.layers = nn.ModuleList([EncoderLayer(embed_dim=embed_dim, n_head=n_head, feedforward_dim=feedforward_dim,
                                                  drop_fc=drop_fc, drop_attn=drop_attn, drop_feed=drop_feed)
                                     for _ in range(n_layers)])

    def forward(self, src_sentence, drop, training=True, src_mask=None, src_key_padding_mask=None):
        """

        :param drop:
        :param training:
        :param src_sentence: 嵌入向量的输入
        :param src_mask: 　　一般只在解码时使用，为了并行一次喂入所有解码部分的输入，所以要用mask来进行掩盖当前时刻之后的位置信息
        :param src_key_padding_mask:　编码部分输入的padding情况，形状为 [batch_size, src_len]
        :return:
        """
        encoder_out = src_sentence
        for layer in self.layers:
            encoder_out = layer(src_sentence, training=training, drop=drop, attn_mask=src_mask,
                                src_key_padding_mask=src_key_padding_mask)
        return encoder_out


class PositionWiseFeedForward(nn.Module):
    def __init__(self, drop_fc=None, dim_model=512, dim_feedforward=2048):
        super(PositionWiseFeedForward, self).__init__()
        self.fc1 = nn.Linear(dim_model, dim_feedforward)
        self.fc2 = nn.Linear(dim_feedforward, dim_model)
        self.relu = nn.ReLU()
        self.drop = nn.Dropout(drop_fc)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.drop(x)
        x = self.fc2(x)
        return x

DecoderLayer解码层

单个解码层的实现

需要与Ｅncoder层进行交互

复制代码 
class DecoderLayer(nn.Module):
    def __init__(self, embed_dim, feedforward_dim, n_head, drop_fc, drop_mask_attn, drop_attn, drop_feed, bias=True):
        """
        解码层：输入嵌入＋位置编码－－＞带掩码的多头注意力－－＞与编码交互的多头注意力－－＞前馈网络
        :param embed_dim:　嵌入维度：512
        :param feedforward_dim:　前馈网络的隐层维度：2048
        :param n_head:　　　多头注意力头数：8　　512//64=8    d_model//batch_size=n_head
        :param drop_fc:　前馈网络中隐层输出的随机丢弃比例
        :param drop_mask_attn:　带掩码的多头注意力层输出
        :param drop_attn:　　　　与编码交互的多头注意力层输出
        :param drop_feed:　　　　前馈网络层输出
        :param bias:　　　　　　　多头注意力层线性输出时候是否加偏置
        """
        super(DecoderLayer, self).__init__()
        # self.out_embedding =
        # TransformerEmbedding(vocab_size=2048, d_model=512, drop=0.1, drop_prob=0.1, max_len=None, device='cpu')
        self.attn = MultiHeadAttention(dim_model=embed_dim, n_head=n_head, bias=bias)
        self.mask_attn = MultiHeadAttention(dim_model=embed_dim, n_head=n_head, bias=bias)
        self.fc = PositionWiseFeedForward(dim_model=embed_dim, drop_fc=drop_fc, dim_feedforward=feedforward_dim)
        # mask_attention层
        self.drop1 = nn.Dropout(drop_mask_attn)
        # encoder_decoder_attention层
        self.drop2 = nn.Dropout(drop_attn)
        # 前馈网络层
        self.drop3 = nn.Dropout(drop_feed)
        self.norm1 = nn.LayerNorm(embed_dim)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.norm3 = nn.LayerNorm(embed_dim)
        self.activation = nn.ReLU()

    def forward(self, tgt, memory, drop=None, tgt_mask=None, memory_mask=None,
                tgt_key_padding_mask=None, memory_key_padding_mask=None):
        """

        :param drop:
        :param memory_mask: 编码器-解码器交互时的注意力掩码，一般为None
        :param tgt_mask:　注意力机制中的掩码矩阵，用于掩盖当前position之后的信息, [tgt_len, tgt_len]
        :param memory: 编码部分的输出（memory）, [src_len,batch_size,embed_dim]
        :param tgt:　解码部分的输入，形状为 [tgt_len,batch_size, embed_dim]
        :param tgt_key_padding_mask:解码部分输入的padding情况，形状为 [batch_size, tgt_len]，填充词的掩码
        :param memory_key_padding_mask: 编码部分输入的padding情况，形状为 [batch_size, src_len]
        :return:
        """
        # out_embedding = self.out_embedding(tgt_sentence)
        # out_embedding = tgt_sentence
        # MultiHeadAttention返回: attn_out_linear, attn_concat, attn_weights
        mask_attn = self.mask_attn(tgt, tgt, tgt, drop_p=drop, training=True,
                                   attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask)[0]
        mask_out = self.norm1(tgt + self.drop1(mask_attn))
        self_attn = self.attn(tgt, key=memory, value=memory, drop_p=drop, training=True,
                              attn_mask=memory_mask, key_padding_mask=memory_key_padding_mask)[0]
        attn_out = self.norm2(mask_out + self.drop2(self_attn))
        attn_out = self.activation(attn_out)
        feed_forward = self.fc(attn_out)
        decoder_out = self.norm3(attn_out+self.drop3(feed_forward))

        return decoder_out

Decoder多层解码块

多个解码层的实现

复制代码 
class Decoder(nn.Module):
    def __init__(self, n_layers, embed_dim, feedforward_dim, n_head, drop_fc,
                 drop_mask_attn, drop_attn, drop_feed, bias=True):
        super(Decoder, self).__init__()
        self.layers = nn.ModuleList([DecoderLayer(embed_dim=embed_dim, feedforward_dim=feedforward_dim, n_head=n_head,
                                                  drop_fc=drop_fc, drop_mask_attn=drop_mask_attn, drop_attn=drop_attn,
                                                  drop_feed=drop_feed, bias=bias) for _ in range(n_layers)])

    def forward(self, tgt_sentence, memory, drop, tgt_mask=None, memory_mask=None, tgt_key_padding_mask=None,
                memory_key_padding_mask=None):
        decoder_out = tgt_sentence
        for layer in self.layers:
            decoder_out = layer(tgt_sentence, memory, drop=drop, tgt_mask=tgt_mask,
                                tgt_key_padding_mask=tgt_key_padding_mask,
                                memory_key_padding_mask=memory_key_padding_mask, memory_mask=None)
        return decoder_out

if __name__ == '__main__':
    src_len = 5
    batch_size = 2
    d_model = 32
    num_head = 1
    n_layer = 6
    src = torch.rand(([batch_size, src_len, d_model]))
    encoder = Decoder(n_layers=n_layer, embed_dim=d_model, feedforward_dim=2048, n_head=num_head, drop_fc=0.1,
                      drop_mask_attn=0.1, drop_attn=0.1, drop_feed=0.1, bias=True)
    print(encoder)

MyTransformer实现完整的Transformer框架

复制代码 
import torch
import torch.nn as nn
import torch.nn.functional as F
from TransformerEmbedding import TransformerEmbedding
from EncoderLayer import Encoder
from DecoderLayer import Decoder


class MyTransformer(nn.Module):
    def __init__(self, src_vocab_size, src_max_len, tgt_vocab_size, tgt_max_len, embed_dim, feedforward_dim, n_head,
                 drop_prob, n_layers):
        super(MyTransformer, self).__init__()
        self.src_embedding = TransformerEmbedding(vocab_size=src_vocab_size, max_len=src_max_len, d_model=embed_dim,
                                                  drop=drop_prob, device='cpu')
        self.tgt_embedding = TransformerEmbedding(vocab_size=tgt_vocab_size, max_len=tgt_max_len, d_model=embed_dim,
                                                  drop=drop_prob, device='cpu')
        self.encoder = Encoder(embed_dim=embed_dim, n_head=n_head, feedforward_dim=feedforward_dim, n_layers=n_layers,
                               drop_fc=drop_prob, drop_attn=drop_prob, drop_feed=drop_prob)
        self.decoder = Decoder(embed_dim=embed_dim, feedforward_dim=feedforward_dim, n_head=n_head, n_layers=n_layers,
                               drop_fc=drop_prob, drop_mask_attn=drop_prob, drop_attn=drop_prob, drop_feed=drop_prob)
        # self.fc = nn.Linear(drop_prob)

    def forward(self, src, tgt, drop_prob, training=True, src_mask=None, tgt_mask=None,
                src_key_padding_mask=None, tgt_key_padding_mask=None):
        src_embedding = self.src_embedding(src)
        tgt_embedding = self.tgt_embedding(tgt)
        memory = self.encoder(src_embedding, drop=drop_prob, training=training,
                              src_key_padding_mask=src_key_padding_mask)
        # print(src_embedding.shape, tgt_embedding.shape, memory.shape)
        print('编码完成－－－－－－－－－－－－－－－－－－－－－－－－－－－－－－')
        decoder = self.decoder(tgt_embedding, memory, drop=drop_prob, tgt_mask=tgt_mask,
                               tgt_key_padding_mask=tgt_key_padding_mask, memory_key_padding_mask=src_key_padding_mask)
        print('解码完成')
        return decoder

    def init_parameters(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

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

if __name__ == '__main__':
    Input = [[1, 2, 3, 4], [5, 6, 2], [1, 7, 8, 9, 10, 11, 12, 13, 14, 15]]
    output = [[1, 1, 2, 3, 4], [5, 5, 6, 2], [1, 7, 7, 8, 9, 10, 11, 12, 13, 14, 15]]
    # 有三个句子［［＇你　好　世　界＇］，　［＇我　很　好＇］，　［＇你　不　是　真　正　的　快　乐＇］］
    # 句子长度不一致进行padding处理
    src_sentence = torch.tensor([])
    for i in range(len(Input)):
        src_sen = torch.tensor([Input[i]])
        src_sen = F.pad(src_sen, (0, len(Input[2]) - len(Input[i])), 'constant', 0)
        src_sentence = torch.cat([src_sentence, src_sen], dim=0)
    # x = src_sentence.unsqueeze(dim=0)
    tgt_sentence = torch.tensor([])
    for i in range(len(output)):
        src_sen = torch.tensor([output[i]])
        src_sen = F.pad(src_sen, (0, len(output[2]) - len(output[i])), 'constant', 0)
        tgt_sentence = torch.cat([tgt_sentence, src_sen], dim=0)
    src_len = 10
    batch_size = 3
    dmodel = 32
    tgt_len = 11
    num_head = 8
    # src = torch.rand((src_len, batch_size, dmodel))  # shape: [src_len, batch_size, embed_dim]
    # src_key_padding_mask = torch.tensor([[True, True, True, False, False],
    #                                      [True, True, True, True, False]])  # shape: [batch_size, src_len]
    #
    # tgt = torch.rand((tgt_len, batch_size, dmodel))  # shape: [tgt_len, batch_size, embed_dim]
    # tgt_key_padding_mask = torch.tensor([[True, True, True, False, False, False],
    #                                      [True, True, True, True, False, False]])  # shape: [batch_size, tgt_len]
    my_transformer = MyTransformer(src_vocab_size=5000, src_max_len=10, tgt_vocab_size=5000, tgt_max_len=11, embed_dim=32,
                                   feedforward_dim=2048, n_head=8, drop_prob=0.1, n_layers=6)
   
    tgt_mask = my_transformer.generate_square_subsequent_mask(tgt_len)
    out = my_transformer(src=src_sentence, tgt=tgt_sentence, drop_prob=0.1, tgt_mask=tgt_mask,
                         src_key_padding_mask=src_sentence,
                         tgt_key_padding_mask=tgt_sentence)
    print(out.shape)
    # torch.Size([3, 11, 32])
相关推荐
美人鱼战士爱学习几秒前
KAG: Boosting LLMs in Professional Domains via Knowledge Augmented Generation
人工智能·集成学习·boosting
有Li2 分钟前
融合先验文本与解剖学知识的多模态回归网络用于舌鳞状细胞癌浸润深度的自动预测|文献速递-文献分享
论文阅读·人工智能·医学生
一水鉴天13 分钟前
整体设计 全面梳理复盘之31 Transformer 九宫格三层架构 Designer 全部功能定稿（初稿）之3
人工智能
速创圈15 分钟前
Sora2 Pro国内接入终极指南：失败退款+无并发限制
人工智能
文心快码 Baidu Comate44 分钟前
双十一将至，用Rules玩转电商场景提效
人工智能·ai编程·文心快码·智能编程助手·comate ai ide
瞻邈1 小时前
LION运行笔记
人工智能·深度学习
CoovallyAIHub1 小时前
外科医生离手术世界模型还有多远？首次提出SurgVeo基准，揭示AI生成手术视频的惊人差距
深度学习·算法·计算机视觉
Serverless 社区1 小时前
助力企业构建 AI 原生应用，函数计算FunctionAI 重塑模型服务与 Agent 全栈生态
大数据·人工智能
大千AI助手1 小时前
参考先验（Reference Priors）详解：理论与Python实践
人工智能·机器学习·贝叶斯·大千ai助手·参考先验·贝叶斯推断·先验
Baihai_IDP1 小时前
面向 LLM 的 GPU 系统工程方法论
人工智能·面试·gpu
热门推荐
01GitHub 镜像站点02UV安装并设置国内源03综合整理：pdf预览显示：你尝试预览的文件可能对你的计算机有害。如果你信任此文件以及其来源，请打开此文件以看其内容，如何解决以正常预览文件04Linux下V2Ray安装配置指南05安娜的档案(Anna’s Archive) 镜像网站/国内最新可访问入口（持续更新）06BongoCat - 跨平台键盘猫动画工具07Labelme从安装到标注：零基础完整指南08《大数据技术原理与应用》实验报告三 熟悉HBase常用操作09智能库存管理的需求预测模型：从业务痛点到落地代码的完整实践10看了下昨日泄露的苹果 App Store 源码……