devbox
盐析白兔
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

Transformer 结构及其代码实现_transformer实现

作者：盐析白兔 | 2024-04-09 11:13:12

赞

踩

transformer实现

一、Transformer 结构图

        如下图，为 Transformer 的整体结构图，左侧为 Transformer Encoder Block，右侧为 Transformer Decoder Block。

        在整体使用中，两个 Block 均被多次重复使用，即上一 Block 的输出向量作为下一 Block 的输入向量。

Transformer Architecture

二、代码实现

        由（一）所介绍，Transformer 是由 TransformerEncoder 和 TransformerDecoder 组成，而这两者又分别是由多个 TransformerEncoderLayers 和 TransformerDecoderLayers 组成（或理解为多个 Block 组成）

        下图代码，建议对照上图内部结构去看，更容易理解一些。

        1.1）TransformerEncoderLayer 代码：

  1. # Transformer Encoder Layer
  2. class TransformerEncoderLayer(Module):
  3.     r"""TransformerEncoderLayer is made up of self-attn and feedforward network.
  4.     This standard encoder layer is based on the paper "Attention Is All You Need".
  5.     Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
  6.     Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in
  7.     Neural Information Processing Systems, pages 6000-6010. Users may modify or implement
  8.     in a different way during application.
  10.     Args:
  11.         d_model: the number of expected features in the input (required).
  12.         nhead: the number of heads in the multiheadattention models (required).
  13.         dim_feedforward: the dimension of the feedforward network model (default=2048).
  14.         dropout: the dropout value (default=0.1).
  15.         activation: the activation function of intermediate layer, relu or gelu (default=relu).
  17.     Examples::
  18.         >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
  19.         >>> src = torch.rand(10, 32, 512)
  20.         >>> out = encoder_layer(src)
  21.     """
  22.  
  23.     def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation="relu"):
  24.         super(TransformerEncoderLayer, self).__init__()
  25.         self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
  26.         # Implementation of Feedforward model
  27.         self.linear1 = Linear(d_model, dim_feedforward)
  28.         self.dropout = Dropout(dropout)
  29.         self.linear2 = Linear(dim_feedforward, d_model)
  30.  
  31.         self.norm1 = LayerNorm(d_model)
  32.         self.norm2 = LayerNorm(d_model)
  33.         self.dropout1 = Dropout(dropout)
  34.         self.dropout2 = Dropout(dropout)
  35.  
  36.         self.activation = _get_activation_fn(activation)
  37.  
  38.     def __setstate__(self, state):
  39.         if 'activation' not in state:
  40.             state['activation'] = F.relu
  41.         super(TransformerEncoderLayer, self).__setstate__(state)
  42.  
  43.     def forward(self, src: Tensor, src_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
  44.         r"""Pass the input through the encoder layer.
  46.         Args:
  47.             src: the sequence to the encoder layer (required).
  48.             src_mask: the mask for the src sequence (optional).
  49.             src_key_padding_mask: the mask for the src keys per batch (optional).
  51.         Shape:
  52.             see the docs in Transformer class.
  53.         """
  54.         # look the picture of transformer encoder
  55.         # Norm(src+Dropout(self_attention(src)))
  56.         src2 = self.self_attn(src, src, src, attn_mask=src_mask,
  57.                               key_padding_mask=src_key_padding_mask)[0]
  58.         src = src + self.dropout1(src2)
  59.         src = self.norm1(src)
  60.  
  61.         # Norm(src+Dropout(Feedforward()))
  62.         src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
  63.         src = src + self.dropout2(src2)
  64.         src = self.norm2(src)
  65.         return src

        1.2）TransformerEncoder 代码（多次执行 TransformerEncoderLayer 里的内容）：

  1. # A stack of N encoder layers
  2. class TransformerEncoder(Module):
  3.     r"""TransformerEncoder is a stack of N encoder layers
  5.     Args:
  6.         encoder_layer: an instance of the TransformerEncoderLayer() class (required).
  7.         num_layers: the number of sub-encoder-layers in the encoder (required).
  8.         norm: the layer normalization component (optional).
  10.     Examples::
  11.         >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
  12.         >>> transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)
  13.         >>> src = torch.rand(10, 32, 512)
  14.         >>> out = transformer_encoder(src)
  15.     """
  16.     __constants__ = ['norm']
  17.  
  18.     def __init__(self, encoder_layer, num_layers, norm=None):
  19.         super(TransformerEncoder, self).__init__()
  20.         self.layers = _get_clones(encoder_layer, num_layers)
  21.         self.num_layers = num_layers
  22.         self.norm = norm
  23.  
  24.     def forward(self, src: Tensor, mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
  25.         r"""Pass the input through the encoder layers in turn.
  27.         Args:
  28.             src: the sequence to the encoder (required).
  29.             mask: the mask for the src sequence (optional).
  30.             src_key_padding_mask: the mask for the src keys per batch (optional).
  32.         Shape:
  33.             see the docs in Transformer class.
  34.         """
  35.         output = src
  36.  
  37.         for mod in self.layers:
  38.             output = mod(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask)
  39.  
  40.         if self.norm is not None:
  41.             output = self.norm(output)
  42.  
  43.         return output

        2.1）TransformerDecoderLayer 代码：

  1. # For language reconstruct
  2. class TransformerDecoderLayer(Module):
  3.     r"""TransformerDecoderLayer is made up of self-attn, multi-head-attn and feedforward network.
  4.     This standard decoder layer is based on the paper "Attention Is All You Need".
  5.     Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
  6.     Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in
  7.     Neural Information Processing Systems, pages 6000-6010. Users may modify or implement
  8.     in a different way during application.
  10.     Args:
  11.         d_model: the number of expected features in the input (required).
  12.         nhead: the number of heads in the multiheadattention models (required).
  13.         dim_feedforward: the dimension of the feedforward network model (default=2048).
  14.         dropout: the dropout value (default=0.1).
  15.         activation: the activation function of intermediate layer, relu or gelu (default=relu).
  17.     Examples::
  18.         >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
  19.         >>> memory = torch.rand(10, 32, 512)
  20.         >>> tgt = torch.rand(20, 32, 512)
  21.         >>> out = decoder_layer(tgt, memory)
  22.     """
  23.     # d_model = 768, nhead = 8
  24.     def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation="relu"):
  25.         super(TransformerDecoderLayer, self).__init__()
  26.         self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
  27.         self.multihead_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
  28.         # Implementation of Feedforward model
  29.         self.linear1 = Linear(d_model, dim_feedforward)
  30.         self.dropout = Dropout(dropout)
  31.         self.linear2 = Linear(dim_feedforward, d_model)
  32.  
  33.         self.norm1 = LayerNorm(d_model)
  34.         self.norm2 = LayerNorm(d_model)
  35.         self.norm3 = LayerNorm(d_model)
  36.         self.dropout1 = Dropout(dropout)
  37.         self.dropout2 = Dropout(dropout)
  38.         self.dropout3 = Dropout(dropout)
  39.  
  40.         self.activation = _get_activation_fn(activation)
  41.  
  42.     def __setstate__(self, state):
  43.         if 'activation' not in state:
  44.             state['activation'] = F.relu
  45.         super(TransformerDecoderLayer, self).__setstate__(state)
  46.  
  47.     # tgt: the sequence to the decoder layer (required).    (20,1,768)
  48.     # memory: the sequence from the last layer of the encoder (required).   (3600,1,768)
  49.     def forward(self, tgt: Tensor, memory: Tensor, tgt_mask: Optional[Tensor] = None, memory_mask: Optional[Tensor] = None,
  50.                 tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:
  51.         r"""Pass the inputs (and mask) through the decoder layer.
  53.         Args:
  54.             tgt: the sequence to the decoder layer (required).
  55.             memory: the sequence from the last layer of the encoder (required).
  56.             tgt_mask: the mask for the tgt sequence (optional).
  57.             memory_mask: the mask for the memory sequence (optional).
  58.             tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
  59.             memory_key_padding_mask: the mask for the memory keys per batch (optional).
  61.         Shape:
  62.             see the docs in Transformer class.
  63.         """ 
  64.  
  65.         # 类比 Transformer Decoder 的结构
  66.         # tgt = Norm(Dropout(attention(tgt,tgt,tgt))+tgt)
  67.         tgt2 = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask,                # Multihead self-attention  tgt2 (20,1,768)
  68.                               key_padding_mask=tgt_key_padding_mask)[0]
  69.         tgt = tgt + self.dropout1(tgt2)                                         # tgt = tgt + dropout1(0.1,tgt2) (20,1,768)   
  70.         tgt = self.norm1(tgt)                                                   # LayerNorm                      (1,768,60,60)
  71.         
  72.         # tgt = Norm(Dropout(attention(tgt,memory,memory))+tgt)
  73.         tgt2 = self.multihead_attn(tgt, memory, memory, attn_mask=memory_mask,  # Multihead self-attention  tgt2 (20,1,768)
  74.                                    key_padding_mask=memory_key_padding_mask)[0]
  75.         tgt = tgt + self.dropout2(tgt2)                                         # (20,1,768)
  76.         tgt = self.norm2(tgt)                                                   # (20,1,768)
  77.         tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))   # tgt2 (20,1,768)
  78.         tgt = tgt + self.dropout3(tgt2)
  79.         tgt = self.norm3(tgt)                   
  80.         return tgt      # (20,1,768)

      2.2）TransformerDecoder 代码（多次执行 TransformerDecoderLayer 里的内容）：

  1. # A stack of N decoder layers
  2. class TransformerDecoder(Module):
  3.     r"""TransformerDecoder is a stack of N decoder layers
  5.     Args:
  6.         decoder_layer: an instance of the TransformerDecoderLayer() class (required).
  7.         num_layers: the number of sub-decoder-layers in the decoder (required).
  8.         norm: the layer normalization component (optional).
  10.     Examples::
  11.         >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
  12.         >>> transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=6)
  13.         >>> memory = torch.rand(10, 32, 512)
  14.         >>> tgt = torch.rand(20, 32, 512)
  15.         >>> out = transformer_decoder(tgt, memory)
  16.     """
  17.     __constants__ = ['norm']
  18.  
  19.     def __init__(self, decoder_layer, num_layers, norm=None):
  20.         super(TransformerDecoder, self).__init__()
  21.         self.layers = _get_clones(decoder_layer, num_layers)
  22.         self.num_layers = num_layers
  23.         self.norm = norm
  24.  
  25.     def forward(self, tgt: Tensor, memory: Tensor, tgt_mask: Optional[Tensor] = None,
  26.                 memory_mask: Optional[Tensor] = None, tgt_key_padding_mask: Optional[Tensor] = None,
  27.                 memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:
  28.         r"""Pass the inputs (and mask) through the decoder layer in turn.
  30.         Args:
  31.             tgt: the sequence to the decoder (required).
  32.             memory: the sequence from the last layer of the encoder (required).
  33.             tgt_mask: the mask for the tgt sequence (optional).
  34.             memory_mask: the mask for the memory sequence (optional).
  35.             tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
  36.             memory_key_padding_mask: the mask for the memory keys per batch (optional).
  38.         Shape:
  39.             see the docs in Transformer class.
  40.         """
  41.         output = tgt
  42.  
  43.         for mod in self.layers:
  44.             output = mod(output, memory, tgt_mask=tgt_mask,
  45.                          memory_mask=memory_mask,
  46.                          tgt_key_padding_mask=tgt_key_padding_mask,
  47.                          memory_key_padding_mask=memory_key_padding_mask)
  48.  
  49.         if self.norm is not None:
  50.             output = self.norm(output)
  51.  
  52.         return output

          2.3）Transformer 代码：

                a、先初始化 TransformerEncoder 和 TransformerDecoder

                b、在 forward( ) 中分别调用他们

  1. # High Architecture of Transformer encoder and Transformer decoder
  2. class Transformer(Module):
  3.     r"""A transformer model. User is able to modify the attributes as needed. The architecture
  4.     is based on the paper "Attention Is All You Need". Ashish Vaswani, Noam Shazeer,
  5.     Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and
  6.     Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information
  7.     Processing Systems, pages 6000-6010. Users can build the BERT(https://arxiv.org/abs/1810.04805)
  8.     model with corresponding parameters.
  10.     Args:
  11.         d_model: the number of expected features in the encoder/decoder inputs (default=512).
  12.         nhead: the number of heads in the multiheadattention models (default=8).
  13.         num_encoder_layers: the number of sub-encoder-layers in the encoder (default=6).
  14.         num_decoder_layers: the number of sub-decoder-layers in the decoder (default=6).
  15.         dim_feedforward: the dimension of the feedforward network model (default=2048).
  16.         dropout: the dropout value (default=0.1).
  17.         activation: the activation function of encoder/decoder intermediate layer, relu or gelu (default=relu).
  18.         custom_encoder: custom encoder (default=None).
  19.         custom_decoder: custom decoder (default=None).
  21.     Examples::
  22.         >>> transformer_model = nn.Transformer(nhead=16, num_encoder_layers=12)
  23.         >>> src = torch.rand((10, 32, 512))
  24.         >>> tgt = torch.rand((20, 32, 512))
  25.         >>> out = transformer_model(src, tgt)
  27.     Note: A full example to apply nn.Transformer module for the word language model is available in
  28.     https://github.com/pytorch/examples/tree/master/word_language_model
  29.     """
  30.  
  31.     def __init__(self, d_model: int = 512, nhead: int = 8, num_encoder_layers: int = 6,
  32.                  num_decoder_layers: int = 6, dim_feedforward: int = 2048, dropout: float = 0.1,
  33.                  activation: str = "relu", custom_encoder: Optional[Any] = None, custom_decoder: Optional[Any] = None) -> None:
  34.         super(Transformer, self).__init__()
  35.  
  36.         if custom_encoder is not None:
  37.             self.encoder = custom_encoder
  38.         else:
  39.             encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout, activation)
  40.             encoder_norm = LayerNorm(d_model)
  41.             self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
  42.  
  43.         if custom_decoder is not None:
  44.             self.decoder = custom_decoder
  45.         else:
  46.             decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout, activation)
  47.             decoder_norm = LayerNorm(d_model)
  48.             self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm)
  49.  
  50.         self._reset_parameters()
  51.  
  52.         self.d_model = d_model
  53.         self.nhead = nhead
  54.  
  55.     def forward(self, src: Tensor, tgt: Tensor, src_mask: Optional[Tensor] = None, tgt_mask: Optional[Tensor] = None,
  56.                 memory_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None,
  57.                 tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:
  58.         r"""Take in and process masked source/target sequences.
  60.         Args:
  61.             src: the sequence to the encoder (required).
  62.             tgt: the sequence to the decoder (required).
  63.             src_mask: the additive mask for the src sequence (optional).
  64.             tgt_mask: the additive mask for the tgt sequence (optional).
  65.             memory_mask: the additive mask for the encoder output (optional).
  66.             src_key_padding_mask: the ByteTensor mask for src keys per batch (optional).
  67.             tgt_key_padding_mask: the ByteTensor mask for tgt keys per batch (optional).
  68.             memory_key_padding_mask: the ByteTensor mask for memory keys per batch (optional).
  70.         Shape:
  71.             - src: :math:`(S, N, E)`.
  72.             - tgt: :math:`(T, N, E)`.
  73.             - src_mask: :math:`(S, S)`.
  74.             - tgt_mask: :math:`(T, T)`.
  75.             - memory_mask: :math:`(T, S)`.
  76.             - src_key_padding_mask: :math:`(N, S)`.
  77.             - tgt_key_padding_mask: :math:`(N, T)`.
  78.             - memory_key_padding_mask: :math:`(N, S)`.
  80.             Note: [src/tgt/memory]_mask ensures that position i is allowed to attend the unmasked
  81.             positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
  82.             while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
  83.             are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
  84.             is provided, it will be added to the attention weight. 
  85.             [src/tgt/memory]_key_padding_mask provides specified elements in the key to be ignored by
  86.             the attention. If a ByteTensor is provided, the non-zero positions will be ignored while the zero
  87.             positions will be unchanged. If a BoolTensor is provided, the positions with the
  88.             value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
  90.             - output: :math:`(T, N, E)`.
  92.             Note: Due to the multi-head attention architecture in the transformer model,
  93.             the output sequence length of a transformer is same as the input sequence
  94.             (i.e. target) length of the decode.
  96.             where S is the source sequence length, T is the target sequence length, N is the
  97.             batch size, E is the feature number
  99.         Examples:
  100.             >>> output = transformer_model(src, tgt, src_mask=src_mask, tgt_mask=tgt_mask)
  101.         """
  102.  
  103.         if src.size(1) != tgt.size(1):
  104.             raise RuntimeError("the batch number of src and tgt must be equal")
  105.  
  106.         if src.size(2) != self.d_model or tgt.size(2) != self.d_model:
  107.             raise RuntimeError("the feature number of src and tgt must be equal to d_model")
  108.  
  109.         memory = self.encoder(src, mask=src_mask, src_key_padding_mask=src_key_padding_mask)
  110.         output = self.decoder(tgt, memory, tgt_mask=tgt_mask, memory_mask=memory_mask,
  111.                               tgt_key_padding_mask=tgt_key_padding_mask,
  112.                               memory_key_padding_mask=memory_key_padding_mask)
  113.         return output
  114.  
  115.     def generate_square_subsequent_mask(self, sz: int) -> Tensor:
  116.         r"""Generate a square mask for the sequence. The masked positions are filled with float('-inf').
  117.             Unmasked positions are filled with float(0.0).
  118.         """
  119.         mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
  120.         mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
  121.         return mask
  122.  
  123.     def _reset_parameters(self):
  124.         r"""Initiate parameters in the transformer model."""
  125.  
  126.         for p in self.parameters():
  127.             if p.dim() > 1:
  128.                 xavier_uniform_(p)

         下面是完整的 transformer.py 文件：

  1. import copy
  2. from typing import Optional, Any
  3.  
  4. import torch
  5. from torch import Tensor
  6. from .. import functional as F
  7. from .module import Module
  8. from .activation import MultiheadAttention
  9. from .container import ModuleList
  10. from ..init import xavier_uniform_
  11. from .dropout import Dropout
  12. from .linear import Linear
  13. from .normalization import LayerNorm
  14.  
  15. # High Architecture of Transformer encoder and Transformer decoder
  16. class Transformer(Module):
  17.     r"""A transformer model. User is able to modify the attributes as needed. The architecture
  18.     is based on the paper "Attention Is All You Need". Ashish Vaswani, Noam Shazeer,
  19.     Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and
  20.     Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information
  21.     Processing Systems, pages 6000-6010. Users can build the BERT(https://arxiv.org/abs/1810.04805)
  22.     model with corresponding parameters.
  24.     Args:
  25.         d_model: the number of expected features in the encoder/decoder inputs (default=512).
  26.         nhead: the number of heads in the multiheadattention models (default=8).
  27.         num_encoder_layers: the number of sub-encoder-layers in the encoder (default=6).
  28.         num_decoder_layers: the number of sub-decoder-layers in the decoder (default=6).
  29.         dim_feedforward: the dimension of the feedforward network model (default=2048).
  30.         dropout: the dropout value (default=0.1).
  31.         activation: the activation function of encoder/decoder intermediate layer, relu or gelu (default=relu).
  32.         custom_encoder: custom encoder (default=None).
  33.         custom_decoder: custom decoder (default=None).
  35.     Examples::
  36.         >>> transformer_model = nn.Transformer(nhead=16, num_encoder_layers=12)
  37.         >>> src = torch.rand((10, 32, 512))
  38.         >>> tgt = torch.rand((20, 32, 512))
  39.         >>> out = transformer_model(src, tgt)
  41.     Note: A full example to apply nn.Transformer module for the word language model is available in
  42.     https://github.com/pytorch/examples/tree/master/word_language_model
  43.     """
  44.  
  45.     def __init__(self, d_model: int = 512, nhead: int = 8, num_encoder_layers: int = 6,
  46.                  num_decoder_layers: int = 6, dim_feedforward: int = 2048, dropout: float = 0.1,
  47.                  activation: str = "relu", custom_encoder: Optional[Any] = None, custom_decoder: Optional[Any] = None) -> None:
  48.         super(Transformer, self).__init__()
  49.  
  50.         if custom_encoder is not None:
  51.             self.encoder = custom_encoder
  52.         else:
  53.             encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout, activation)
  54.             encoder_norm = LayerNorm(d_model)
  55.             self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
  56.  
  57.         if custom_decoder is not None:
  58.             self.decoder = custom_decoder
  59.         else:
  60.             decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout, activation)
  61.             decoder_norm = LayerNorm(d_model)
  62.             self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm)
  63.  
  64.         self._reset_parameters()
  65.  
  66.         self.d_model = d_model
  67.         self.nhead = nhead
  68.  
  69.     def forward(self, src: Tensor, tgt: Tensor, src_mask: Optional[Tensor] = None, tgt_mask: Optional[Tensor] = None,
  70.                 memory_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None,
  71.                 tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:
  72.         r"""Take in and process masked source/target sequences.
  74.         Args:
  75.             src: the sequence to the encoder (required).
  76.             tgt: the sequence to the decoder (required).
  77.             src_mask: the additive mask for the src sequence (optional).
  78.             tgt_mask: the additive mask for the tgt sequence (optional).
  79.             memory_mask: the additive mask for the encoder output (optional).
  80.             src_key_padding_mask: the ByteTensor mask for src keys per batch (optional).
  81.             tgt_key_padding_mask: the ByteTensor mask for tgt keys per batch (optional).
  82.             memory_key_padding_mask: the ByteTensor mask for memory keys per batch (optional).
  84.         Shape:
  85.             - src: :math:`(S, N, E)`.
  86.             - tgt: :math:`(T, N, E)`.
  87.             - src_mask: :math:`(S, S)`.
  88.             - tgt_mask: :math:`(T, T)`.
  89.             - memory_mask: :math:`(T, S)`.
  90.             - src_key_padding_mask: :math:`(N, S)`.
  91.             - tgt_key_padding_mask: :math:`(N, T)`.
  92.             - memory_key_padding_mask: :math:`(N, S)`.
  94.             Note: [src/tgt/memory]_mask ensures that position i is allowed to attend the unmasked
  95.             positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
  96.             while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
  97.             are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
  98.             is provided, it will be added to the attention weight. 
  99.             [src/tgt/memory]_key_padding_mask provides specified elements in the key to be ignored by
  100.             the attention. If a ByteTensor is provided, the non-zero positions will be ignored while the zero
  101.             positions will be unchanged. If a BoolTensor is provided, the positions with the
  102.             value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
  104.             - output: :math:`(T, N, E)`.
  106.             Note: Due to the multi-head attention architecture in the transformer model,
  107.             the output sequence length of a transformer is same as the input sequence
  108.             (i.e. target) length of the decode.
  110.             where S is the source sequence length, T is the target sequence length, N is the
  111.             batch size, E is the feature number
  113.         Examples:
  114.             >>> output = transformer_model(src, tgt, src_mask=src_mask, tgt_mask=tgt_mask)
  115.         """
  116.  
  117.         if src.size(1) != tgt.size(1):
  118.             raise RuntimeError("the batch number of src and tgt must be equal")
  119.  
  120.         if src.size(2) != self.d_model or tgt.size(2) != self.d_model:
  121.             raise RuntimeError("the feature number of src and tgt must be equal to d_model")
  122.  
  123.         memory = self.encoder(src, mask=src_mask, src_key_padding_mask=src_key_padding_mask)
  124.         output = self.decoder(tgt, memory, tgt_mask=tgt_mask, memory_mask=memory_mask,
  125.                               tgt_key_padding_mask=tgt_key_padding_mask,
  126.                               memory_key_padding_mask=memory_key_padding_mask)
  127.         return output
  128.  
  129.     def generate_square_subsequent_mask(self, sz: int) -> Tensor:
  130.         r"""Generate a square mask for the sequence. The masked positions are filled with float('-inf').
  131.             Unmasked positions are filled with float(0.0).
  132.         """
  133.         mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
  134.         mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
  135.         return mask
  136.  
  137.     def _reset_parameters(self):
  138.         r"""Initiate parameters in the transformer model."""
  139.  
  140.         for p in self.parameters():
  141.             if p.dim() > 1:
  142.                 xavier_uniform_(p)
  143.  
  144. # A stack of N encoder layers
  145. class TransformerEncoder(Module):
  146.     r"""TransformerEncoder is a stack of N encoder layers
  148.     Args:
  149.         encoder_layer: an instance of the TransformerEncoderLayer() class (required).
  150.         num_layers: the number of sub-encoder-layers in the encoder (required).
  151.         norm: the layer normalization component (optional).
  153.     Examples::
  154.         >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
  155.         >>> transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)
  156.         >>> src = torch.rand(10, 32, 512)
  157.         >>> out = transformer_encoder(src)
  158.     """
  159.     __constants__ = ['norm']
  160.  
  161.     def __init__(self, encoder_layer, num_layers, norm=None):
  162.         super(TransformerEncoder, self).__init__()
  163.         self.layers = _get_clones(encoder_layer, num_layers)
  164.         self.num_layers = num_layers
  165.         self.norm = norm
  166.  
  167.     def forward(self, src: Tensor, mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
  168.         r"""Pass the input through the encoder layers in turn.
  170.         Args:
  171.             src: the sequence to the encoder (required).
  172.             mask: the mask for the src sequence (optional).
  173.             src_key_padding_mask: the mask for the src keys per batch (optional).
  175.         Shape:
  176.             see the docs in Transformer class.
  177.         """
  178.         output = src
  179.  
  180.         for mod in self.layers:
  181.             output = mod(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask)
  182.  
  183.         if self.norm is not None:
  184.             output = self.norm(output)
  185.  
  186.         return output
  187.  
  188. # A stack of N decoder layers
  189. class TransformerDecoder(Module):
  190.     r"""TransformerDecoder is a stack of N decoder layers
  192.     Args:
  193.         decoder_layer: an instance of the TransformerDecoderLayer() class (required).
  194.         num_layers: the number of sub-decoder-layers in the decoder (required).
  195.         norm: the layer normalization component (optional).
  197.     Examples::
  198.         >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
  199.         >>> transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=6)
  200.         >>> memory = torch.rand(10, 32, 512)
  201.         >>> tgt = torch.rand(20, 32, 512)
  202.         >>> out = transformer_decoder(tgt, memory)
  203.     """
  204.     __constants__ = ['norm']
  205.  
  206.     def __init__(self, decoder_layer, num_layers, norm=None):
  207.         super(TransformerDecoder, self).__init__()
  208.         self.layers = _get_clones(decoder_layer, num_layers)
  209.         self.num_layers = num_layers
  210.         self.norm = norm
  211.  
  212.     def forward(self, tgt: Tensor, memory: Tensor, tgt_mask: Optional[Tensor] = None,
  213.                 memory_mask: Optional[Tensor] = None, tgt_key_padding_mask: Optional[Tensor] = None,
  214.                 memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:
  215.         r"""Pass the inputs (and mask) through the decoder layer in turn.
  217.         Args:
  218.             tgt: the sequence to the decoder (required).
  219.             memory: the sequence from the last layer of the encoder (required).
  220.             tgt_mask: the mask for the tgt sequence (optional).
  221.             memory_mask: the mask for the memory sequence (optional).
  222.             tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
  223.             memory_key_padding_mask: the mask for the memory keys per batch (optional).
  225.         Shape:
  226.             see the docs in Transformer class.
  227.         """
  228.         output = tgt
  229.  
  230.         for mod in self.layers:
  231.             output = mod(output, memory, tgt_mask=tgt_mask,
  232.                          memory_mask=memory_mask,
  233.                          tgt_key_padding_mask=tgt_key_padding_mask,
  234.                          memory_key_padding_mask=memory_key_padding_mask)
  235.  
  236.         if self.norm is not None:
  237.             output = self.norm(output)
  238.  
  239.         return output
  240.  
  241. # Transformer Encoder Layer
  242. class TransformerEncoderLayer(Module):
  243.     r"""TransformerEncoderLayer is made up of self-attn and feedforward network.
  244.     This standard encoder layer is based on the paper "Attention Is All You Need".
  245.     Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
  246.     Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in
  247.     Neural Information Processing Systems, pages 6000-6010. Users may modify or implement
  248.     in a different way during application.
  250.     Args:
  251.         d_model: the number of expected features in the input (required).
  252.         nhead: the number of heads in the multiheadattention models (required).
  253.         dim_feedforward: the dimension of the feedforward network model (default=2048).
  254.         dropout: the dropout value (default=0.1).
  255.         activation: the activation function of intermediate layer, relu or gelu (default=relu).
  257.     Examples::
  258.         >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
  259.         >>> src = torch.rand(10, 32, 512)
  260.         >>> out = encoder_layer(src)
  261.     """
  262.  
  263.     def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation="relu"):
  264.         super(TransformerEncoderLayer, self).__init__()
  265.         self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
  266.         # Implementation of Feedforward model
  267.         self.linear1 = Linear(d_model, dim_feedforward)
  268.         self.dropout = Dropout(dropout)
  269.         self.linear2 = Linear(dim_feedforward, d_model)
  270.  
  271.         self.norm1 = LayerNorm(d_model)
  272.         self.norm2 = LayerNorm(d_model)
  273.         self.dropout1 = Dropout(dropout)
  274.         self.dropout2 = Dropout(dropout)
  275.  
  276.         self.activation = _get_activation_fn(activation)
  277.  
  278.     def __setstate__(self, state):
  279.         if 'activation' not in state:
  280.             state['activation'] = F.relu
  281.         super(TransformerEncoderLayer, self).__setstate__(state)
  282.  
  283.     def forward(self, src: Tensor, src_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
  284.         r"""Pass the input through the encoder layer.
  286.         Args:
  287.             src: the sequence to the encoder layer (required).
  288.             src_mask: the mask for the src sequence (optional).
  289.             src_key_padding_mask: the mask for the src keys per batch (optional).
  291.         Shape:
  292.             see the docs in Transformer class.
  293.         """
  294.         # look the picture of transformer encoder
  295.         # Norm(src+Dropout(self_attention(src)))
  296.         src2 = self.self_attn(src, src, src, attn_mask=src_mask,
  297.                               key_padding_mask=src_key_padding_mask)[0]
  298.         src = src + self.dropout1(src2)
  299.         src = self.norm1(src)
  300.  
  301.         # Norm(src+Dropout(Feedforward()))
  302.         src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
  303.         src = src + self.dropout2(src2)
  304.         src = self.norm2(src)
  305.         return src
  306.  
  307. # For language reconstruct
  308. class TransformerDecoderLayer(Module):
  309.     r"""TransformerDecoderLayer is made up of self-attn, multi-head-attn and feedforward network.
  310.     This standard decoder layer is based on the paper "Attention Is All You Need".
  311.     Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
  312.     Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in
  313.     Neural Information Processing Systems, pages 6000-6010. Users may modify or implement
  314.     in a different way during application.
  316.     Args:
  317.         d_model: the number of expected features in the input (required).
  318.         nhead: the number of heads in the multiheadattention models (required).
  319.         dim_feedforward: the dimension of the feedforward network model (default=2048).
  320.         dropout: the dropout value (default=0.1).
  321.         activation: the activation function of intermediate layer, relu or gelu (default=relu).
  323.     Examples::
  324.         >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
  325.         >>> memory = torch.rand(10, 32, 512)
  326.         >>> tgt = torch.rand(20, 32, 512)
  327.         >>> out = decoder_layer(tgt, memory)
  328.     """
  329.     # d_model = 768, nhead = 8
  330.     def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation="relu"):
  331.         super(TransformerDecoderLayer, self).__init__()
  332.         self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
  333.         self.multihead_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
  334.         # Implementation of Feedforward model
  335.         self.linear1 = Linear(d_model, dim_feedforward)
  336.         self.dropout = Dropout(dropout)
  337.         self.linear2 = Linear(dim_feedforward, d_model)
  338.  
  339.         self.norm1 = LayerNorm(d_model)
  340.         self.norm2 = LayerNorm(d_model)
  341.         self.norm3 = LayerNorm(d_model)
  342.         self.dropout1 = Dropout(dropout)
  343.         self.dropout2 = Dropout(dropout)
  344.         self.dropout3 = Dropout(dropout)
  345.  
  346.         self.activation = _get_activation_fn(activation)
  347.  
  348.     def __setstate__(self, state):
  349.         if 'activation' not in state:
  350.             state['activation'] = F.relu
  351.         super(TransformerDecoderLayer, self).__setstate__(state)
  352.  
  353.     # tgt: the sequence to the decoder layer (required).    (20,1,768)
  354.     # memory: the sequence from the last layer of the encoder (required).   (3600,1,768)
  355.     def forward(self, tgt: Tensor, memory: Tensor, tgt_mask: Optional[Tensor] = None, memory_mask: Optional[Tensor] = None,
  356.                 tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:
  357.         r"""Pass the inputs (and mask) through the decoder layer.
  359.         Args:
  360.             tgt: the sequence to the decoder layer (required).
  361.             memory: the sequence from the last layer of the encoder (required).
  362.             tgt_mask: the mask for the tgt sequence (optional).
  363.             memory_mask: the mask for the memory sequence (optional).
  364.             tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
  365.             memory_key_padding_mask: the mask for the memory keys per batch (optional).
  367.         Shape:
  368.             see the docs in Transformer class.
  369.         """ 
  370.  
  371.         # 类比 Transformer Decoder 的结构
  372.         # tgt = Norm(Dropout(attention(tgt,tgt,tgt))+tgt)
  373.         tgt2 = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask,                # Multihead self-attention  tgt2 (20,1,768)
  374.                               key_padding_mask=tgt_key_padding_mask)[0]
  375.         tgt = tgt + self.dropout1(tgt2)                                         # tgt = tgt + dropout1(0.1,tgt2) (20,1,768)   
  376.         tgt = self.norm1(tgt)                                                   # LayerNorm                      (1,768,60,60)
  377.         
  378.         # tgt = Norm(Dropout(attention(tgt,memory,memory))+tgt)
  379.         tgt2 = self.multihead_attn(tgt, memory, memory, attn_mask=memory_mask,  # Multihead self-attention  tgt2 (20,1,768)
  380.                                    key_padding_mask=memory_key_padding_mask)[0]
  381.         tgt = tgt + self.dropout2(tgt2)                                         # (20,1,768)
  382.         tgt = self.norm2(tgt)                                                   # (20,1,768)
  383.         tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))   # tgt2 (20,1,768)
  384.         tgt = tgt + self.dropout3(tgt2)
  385.         tgt = self.norm3(tgt)                   
  386.         return tgt      # (20,1,768)
  387.  
  388.  
  389. def _get_clones(module, N):
  390.     return ModuleList([copy.deepcopy(module) for i in range(N)])
  391.  
  392.  
  393. def _get_activation_fn(activation):
  394.     if activation == "relu":
  395.         return F.relu
  396.     elif activation == "gelu":
  397.         return F.gelu
  398.  
  399.     raise RuntimeError("activation should be relu/gelu, not {}".format(activation))

 三、Transformer 中的多头注意力机制

        Transformer 中多次使用了多头注意力机制。

        在 EncoderLayer 中，使用了一次多头自注意力机制。

        在 DecoderLayer 中，先使用了一次多头自注意力机制，紧接着使用了一次多头非自注意力机制（k 为 tgt，q、v 为memory，是从上一个 encoder block 输出的结果）。

        注意力机制的代码实现如下：

  1. # 多头注意力机制
  2. class MultiheadAttention(Module):
  3.     r"""Allows the model to jointly attend to information
  4.     from different representation subspaces.
  5.     See reference: Attention Is All You Need
  7.     .. math::
  8.         \text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O
  9.         \text{where} head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)
  11.     Args:
  12.         embed_dim: total dimension of the model.
  13.         num_heads: parallel attention heads.
  14.         dropout: a Dropout layer on attn_output_weights. Default: 0.0.
  15.         bias: add bias as module parameter. Default: True.
  16.         add_bias_kv: add bias to the key and value sequences at dim=0.
  17.         add_zero_attn: add a new batch of zeros to the key and
  18.                        value sequences at dim=1.
  19.         kdim: total number of features in key. Default: None.
  20.         vdim: total number of features in value. Default: None.
  22.         Note: if kdim and vdim are None, they will be set to embed_dim such that
  23.         query, key, and value have the same number of features.
  25.     Examples::
  27.         >>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
  28.         >>> attn_output, attn_output_weights = multihead_attn(query, key, value)
  29.     """
  30.     bias_k: Optional[torch.Tensor]
  31.     bias_v: Optional[torch.Tensor]
  32.  
  33.     def __init__(self, embed_dim, num_heads, dropout=0., bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None):
  34.         super(MultiheadAttention, self).__init__()
  35.         self.embed_dim = embed_dim
  36.         self.kdim = kdim if kdim is not None else embed_dim
  37.         self.vdim = vdim if vdim is not None else embed_dim
  38.         self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim
  39.  
  40.         self.num_heads = num_heads
  41.         self.dropout = dropout
  42.         self.head_dim = embed_dim // num_heads
  43.         assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
  44.  
  45.         if self._qkv_same_embed_dim is False:
  46.             self.q_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))
  47.             self.k_proj_weight = Parameter(torch.Tensor(embed_dim, self.kdim))
  48.             self.v_proj_weight = Parameter(torch.Tensor(embed_dim, self.vdim))
  49.             self.register_parameter('in_proj_weight', None)
  50.         else:
  51.             self.in_proj_weight = Parameter(torch.empty(3 * embed_dim, embed_dim))
  52.             self.register_parameter('q_proj_weight', None)
  53.             self.register_parameter('k_proj_weight', None)
  54.             self.register_parameter('v_proj_weight', None)
  55.  
  56.         if bias:
  57.             self.in_proj_bias = Parameter(torch.empty(3 * embed_dim))
  58.         else:
  59.             self.register_parameter('in_proj_bias', None)
  60.         self.out_proj = _LinearWithBias(embed_dim, embed_dim)
  61.  
  62.         if add_bias_kv:
  63.             self.bias_k = Parameter(torch.empty(1, 1, embed_dim))
  64.             self.bias_v = Parameter(torch.empty(1, 1, embed_dim))
  65.         else:
  66.             self.bias_k = self.bias_v = None
  67.  
  68.         self.add_zero_attn = add_zero_attn
  69.  
  70.         self._reset_parameters()
  71.  
  72.     def _reset_parameters(self):
  73.         if self._qkv_same_embed_dim:
  74.             xavier_uniform_(self.in_proj_weight)
  75.         else:
  76.             xavier_uniform_(self.q_proj_weight)
  77.             xavier_uniform_(self.k_proj_weight)
  78.             xavier_uniform_(self.v_proj_weight)
  79.  
  80.         if self.in_proj_bias is not None:
  81.             constant_(self.in_proj_bias, 0.)
  82.             constant_(self.out_proj.bias, 0.)
  83.         if self.bias_k is not None:
  84.             xavier_normal_(self.bias_k)
  85.         if self.bias_v is not None:
  86.             xavier_normal_(self.bias_v)
  87.  
  88.     def __setstate__(self, state):
  89.         # Support loading old MultiheadAttention checkpoints generated by v1.1.0
  90.         if '_qkv_same_embed_dim' not in state:
  91.             state['_qkv_same_embed_dim'] = True
  92.  
  93.         super(MultiheadAttention, self).__setstate__(state)
  94.  
  95.     def forward(self, query, key, value, key_padding_mask=None,
  96.                 need_weights=True, attn_mask=None):
  97.         # type: (Tensor, Tensor, Tensor, Optional[Tensor], bool, Optional[Tensor]) -> Tuple[Tensor, Optional[Tensor]]
  98.         r"""
  99.     Args:
  100.         query, key, value: map a query and a set of key-value pairs to an output.
  101.             See "Attention Is All You Need" for more details.
  102.         key_padding_mask: if provided, specified padding elements in the key will
  103.             be ignored by the attention. When given a binary mask and a value is True,
  104.             the corresponding value on the attention layer will be ignored. When given
  105.             a byte mask and a value is non-zero, the corresponding value on the attention
  106.             layer will be ignored
  107.         need_weights: output attn_output_weights.
  108.         attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
  109.             the batches while a 3D mask allows to specify a different mask for the entries of each batch.
  111.     Shape:
  112.         - Inputs:
  113.         - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
  114.           the embedding dimension.
  115.         - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
  116.           the embedding dimension.
  117.         - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
  118.           the embedding dimension.
  119.         - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
  120.           If a ByteTensor is provided, the non-zero positions will be ignored while the position
  121.           with the zero positions will be unchanged. If a BoolTensor is provided, the positions with the
  122.           value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
  123.         - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
  124.           3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,
  125.           S is the source sequence length. attn_mask ensure that position i is allowed to attend the unmasked
  126.           positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
  127.           while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
  128.           is not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
  129.           is provided, it will be added to the attention weight.
  131.         - Outputs:
  132.         - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
  133.           E is the embedding dimension.
  134.         - attn_output_weights: :math:`(N, L, S)` where N is the batch size,
  135.           L is the target sequence length, S is the source sequence length.
  136.         """
  137.         if not self._qkv_same_embed_dim:
  138.             return F.multi_head_attention_forward(
  139.                 query, key, value, self.embed_dim, self.num_heads,
  140.                 self.in_proj_weight, self.in_proj_bias,
  141.                 self.bias_k, self.bias_v, self.add_zero_attn,
  142.                 self.dropout, self.out_proj.weight, self.out_proj.bias,
  143.                 training=self.training,
  144.                 key_padding_mask=key_padding_mask, need_weights=need_weights,
  145.                 attn_mask=attn_mask, use_separate_proj_weight=True,
  146.                 q_proj_weight=self.q_proj_weight, k_proj_weight=self.k_proj_weight,
  147.                 v_proj_weight=self.v_proj_weight)
  148.         else:
  149.             return F.multi_head_attention_forward(
  150.                 query, key, value, self.embed_dim, self.num_heads,
  151.                 self.in_proj_weight, self.in_proj_bias,
  152.                 self.bias_k, self.bias_v, self.add_zero_attn,
  153.                 self.dropout, self.out_proj.weight, self.out_proj.bias,
  154.                 training=self.training,
  155.                 key_padding_mask=key_padding_mask, need_weights=need_weights,
  156.                 attn_mask=attn_mask)

         此博客为个人学习笔记，如有错误，欢迎指正！感谢各位大佬！

        

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/盐析白兔/article/detail/392360
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号