解析Torch中`Transformer`

解析torch官方代码脚本文件：transformer.py。版本：1.9.1+cu111。

首先查看《Torch中多头注意力MultiheadAttention的中文注释》解析；
最后查看下方transformer解析。
话不多说，看代码吧！

import copy
from typing import Optional, Any

import torch
from torch import Tensor
from .. import functional as F
from .module import Module
from .activation import MultiheadAttention
from .container import ModuleList
from ..init import xavier_uniform_
from .dropout import Dropout
from .linear import Linear
from .normalization import LayerNorm


class Transformer(Module):
    r"""这是一个变换器模型，用户可以根据需要修改其属性。该架构基于论文《Attention Is All You Need》。
    该论文由Ashish Vaswani、Noam Shazeer、Niki Parmar、Jakob Uszkoreit、Llion Jones、Aidan N Gomez、Lukasz Kaiser和Illia Polosukhin于2017年发表，
    在神经信息处理系统进展（Advances in Neural Information Processing Systems）的第6000至6010页。用户可以使用对应的参数构建BERT模型（参见https://arxiv.org/abs/1810.04805）。

    参数如下：

        d_model：编码器/解码器输入中预期的特征数量（默认为512）。
        nhead：多头注意力模型中的头数（默认为8）。
        num_encoder_layers：编码器中子编码层的数量（默认为6）。
        num_decoder_layers：解码器中子解码层的数量（默认为6）。
        dim_feedforward：前馈网络模型的维度（默认为2048）。
        dropout：丢弃率（默认为0.1）。
        activation：编码器/解码器中间层的激活函数，可选relu或gelu（默认为relu）。
        custom_encoder：自定义编码器（默认为None）。
        custom_decoder：自定义解码器（默认为None）。
        layer_norm_eps：层归一化组件中的eps值（默认为1e-5）。
        batch_first：如果设为True，则输入和输出张量将以（批次，序列，特征）的形式提供。默认值为False（序列，批次，特征）。

    Examples::
        >>> transformer_model = nn.Transformer(nhead=16, num_encoder_layers=12)
        >>> src = torch.rand((10, 32, 512))
        >>> tgt = torch.rand((20, 32, 512))
        >>> out = transformer_model(src, tgt)

    Note: A full example to apply nn.Transformer module for the word language model is available in
    https://github.com/pytorch/examples/tree/master/word_language_model

    __init__：
            batch_first: 如果为 True，则输入和输出张量将按 (batch, seq, feature) 的顺序提供；否则，顺序为 (seq, batch, feature)。
        编码器和解码器组件:
            self.encoder: 如果提供了 custom_encoder，则直接使用这个自定义编码器；
                否则，创建一个标准的 TransformerEncoder 实例，它由多个 TransformerEncoderLayer 层堆叠构成。
            self.decoder: 同样地，如果提供了 custom_decoder，则直接使用；
                否则，创建一个标准的 TransformerDecoder 实例，由多个 TransformerDecoderLayer 层构成。
    """

    def __init__(self, d_model: int = 512, nhead: int = 8, num_encoder_layers: int = 6,
                 num_decoder_layers: int = 6, dim_feedforward: int = 2048, dropout: float = 0.1,
                 activation: str = "relu", custom_encoder: Optional[Any] = None, custom_decoder: Optional[Any] = None,
                 layer_norm_eps: float = 1e-5, batch_first: bool = False,
                 device=None, dtype=None) -> None:
        factory_kwargs = {'device': device, 'dtype': dtype}
        super(Transformer, self).__init__()

        if custom_encoder is not None:
            self.encoder = custom_encoder
        else:
            encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout,
                                                    activation, layer_norm_eps, batch_first,
                                                    **factory_kwargs)
            encoder_norm = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
            self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)

        if custom_decoder is not None:
            self.decoder = custom_decoder
        else:
            decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout,
                                                    activation, layer_norm_eps, batch_first,
                                                    **factory_kwargs)
            decoder_norm = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
            self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm)

        self._reset_parameters()

        self.d_model = d_model
        self.nhead = nhead

        self.batch_first = batch_first

    def forward(self, src: Tensor, tgt: Tensor, src_mask: Optional[Tensor] = None, tgt_mask: Optional[Tensor] = None,
                memory_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None,
                tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:
        r"""接受并处理带有掩码的源序列和目标序列。

        参数:
            src: 输入到编码器的序列（必需）。
            tgt: 输入到解码器的序列（必需）。
            src_mask: 源序列的加性掩码（可选）。
            tgt_mask: 目标序列的加性掩码（可选）。
            memory_mask: 编码器输出的加性掩码（可选）。
            src_key_padding_mask: 每批源键的ByteTensor掩码（可选）。
            tgt_key_padding_mask: 每批目标键的ByteTensor掩码（可选）。
            memory_key_padding_mask: 每批记忆键的ByteTensor掩码（可选）。

        形状:
            - src: :math:`(S, N, E)`，如果`batch_first`为真，则为`(N, S, E)`。
            - tgt: :math:`(T, N, E)`，如果`batch_first`为真，则为`(N, T, E)`。
            - src_mask: :math:`(S, S)`。
            - tgt_mask: :math:`(T, T)`。
            - memory_mask: :math:`(T, S)`。
            - src_key_padding_mask: :math:`(N, S)`。
            - tgt_key_padding_mask: :math:`(N, T)`。
            - memory_key_padding_mask: :math:`(N, S)`。

        注意: [src/tgt/memory]_mask确保位置i能访问未被掩码的位置。如果提供的是ByteTensor，非零位置不允许访问，而零位置保持不变。如果提供的是BoolTensor，`True`的位置不允许访问，而`False`值将保持不变。如果提供的是FloatTensor，它将被添加到注意力权重中。
        [src/tgt/memory]_key_padding_mask提供了在键中应被注意力忽略的指定元素。如果提供的是ByteTensor，非零位置将被忽略，而零位置保持不变。如果提供的是BoolTensor，值为`True`的位置将被忽略，而值为`False`的位置将保持不变。

        - 输出: :math:`(T, N, E)`，如果`batch_first`为真，则为`(N, T, E)`。

        注意: 由于变换器模型中的多头注意力架构，变换器的输出序列长度与输入序列（即解码的目标）的长度相同。

        其中S是源序列长度，T是目标序列长度，N是批次大小，E是特征数量。


        Examples:
            >>> output = transformer_model(src, tgt, src_mask=src_mask, tgt_mask=tgt_mask)
        """
        if not self.batch_first and src.size(1) != tgt.size(1):
            raise RuntimeError("the batch number of src and tgt must be equal")
        elif self.batch_first and src.size(0) != tgt.size(0):
            raise RuntimeError("the batch number of src and tgt must be equal")

        if src.size(2) != self.d_model or tgt.size(2) != self.d_model:
            raise RuntimeError("the feature number of src and tgt must be equal to d_model")

        memory = self.encoder(src, mask=src_mask, src_key_padding_mask=src_key_padding_mask)
        output = self.decoder(tgt, memory, tgt_mask=tgt_mask, memory_mask=memory_mask,
                              tgt_key_padding_mask=tgt_key_padding_mask,
                              memory_key_padding_mask=memory_key_padding_mask)
        return output

    def generate_square_subsequent_mask(self, sz: int) -> Tensor:
        r"""Generate a square mask for the sequence. The masked positions are filled with float('-inf').
            Unmasked positions are filled with float(0.0).
        """
        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
        return mask

    def _reset_parameters(self):
        r"""Initiate parameters in the transformer model."""

        for p in self.parameters():
            if p.dim() > 1:
                xavier_uniform_(p)


class TransformerEncoder(Module):
    r"""TransformerEncoder is a stack of N encoder layers

    参数:
    - encoder_layer: TransformerEncoderLayer类的一个实例（必需）。
    - num_layers: 编码器中子编码器层的数量（必需）。
    - norm: 层归一化组件（可选）。

    __init__：
        encoder_layer: 这是单个编码器层的实例，通常由 TransformerEncoderLayer 构造。
        num_layers: 表示要堆叠的编码器层数量。
        norm: 可选参数，用于指定层归一化组件。
        在初始化过程中，通过 _get_clones 函数复制 encoder_layer 指定次数来创建编码器层的列表，这个列表存储在 self.layers 中。self.num_layers 存储了编码器层数量，而 self.norm 则保存了提供的层归一化组件（如果有）

    Examples::
        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
        >>> transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)
        >>> src = torch.rand(10, 32, 512)
        >>> out = transformer_encoder(src)
    """
    __constants__ = ['norm']

    def __init__(self, encoder_layer, num_layers, norm=None):
        super(TransformerEncoder, self).__init__()
        self.layers = _get_clones(encoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm

    def forward(self, src: Tensor, mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
        r"""依次通过编码器层传递输入.

        参数:
        Src:给编码器的序列(必需的)。
        Mask: SRC序列的掩码(可选)。
        Src_key_padding_mask:每批SRC键的掩码(可选)。

        形状:
        请参阅Transformer类中的文档。

        流程：
            src: 输入序列，是必须提供的参数。
            mask: 可选参数，用于输入序列的掩码。
            src_key_padding_mask: 可选参数，用于每批次源键的掩码。
            在前向传播中，数据 src 依次通过每一个编码器层，其中可以使用 mask 和 src_key_padding_mask 来控制哪些部分需要被忽略。最后，如果提供了层归一化组件 norm，那么整个编码器的输出将会通过该组件进行归一化处理
        """
        output = src

        for mod in self.layers:
            output = mod(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask)

        if self.norm is not None:
            output = self.norm(output)

        return output


class TransformerDecoder(Module):
    r"""TransformerDecoder 是由 N 个解码器层堆叠而成的模块

    参数:
        decoder_layer: 必须提供的 TransformerDecoderLayer() 类的一个实例。这定义了单个解码器层的行为和结构。
        num_layers: 解码器中包含的子解码器层数量，这是一个必需的参数，决定了解码器的深度。
        norm: 层归一化组件，这是一个可选参数，用于在解码器所有层的输出之后执行归一化操作，有助于稳定训练过程并加速收敛。

    Examples::
        >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
        >>> transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=6)
        >>> memory = torch.rand(10, 32, 512)
        >>> tgt = torch.rand(20, 32, 512)
        >>> out = transformer_decoder(tgt, memory)
    """
    __constants__ = ['norm']

    def __init__(self, decoder_layer, num_layers, norm=None):
        super(TransformerDecoder, self).__init__()
        self.layers = _get_clones(decoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm

    def forward(self, tgt: Tensor, memory: Tensor, tgt_mask: Optional[Tensor] = None,
                memory_mask: Optional[Tensor] = None, tgt_key_padding_mask: Optional[Tensor] = None,
                memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:
        r"""依次通过解码器层传递输入(和掩码)。

        参数:
            Tgt:到解码器的序列(必需)。
            存储器:来自编码器最后一层的序列(必需)。
            Tgt_mask: TGT序列的掩码(可选)。
            Memory_mask:内存序列的掩码(可选)。
            Tgt_key_padding_mask:每批TGT键的掩码(可选)。
            Memory_key_padding_mask:每批内存键的掩码(可选)。

        形状:
            请参阅Transformer类中的文档。

        """
        output = tgt

        for mod in self.layers:
            output = mod(output, memory, tgt_mask=tgt_mask,
                         memory_mask=memory_mask,
                         tgt_key_padding_mask=tgt_key_padding_mask,
                         memory_key_padding_mask=memory_key_padding_mask)

        if self.norm is not None:
            output = self.norm(output)

        return output

class TransformerEncoderLayer(Module):
    r"""`TransformerEncoderLayer` 由自注意力(self-attn)和前馈网络(feedforward network)组成。
    这个标准的编码器层基于论文 "Attention Is All You Need"。
    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
    Lukasz Kaiser, 和 Illia Polosukhin。2017年。注意力就是你所需要的。在《神经信息处理系统进展》期刊中，
    第6000至6010页。用户在应用过程中可以对其进行修改或以不同方式实现。

    参数:
        d_model: 输入中预期特征的数量（必需）。
        nhead: 多头注意力模型中的头数（必需）。
        dim_feedforward: 前馈网络模型的维度（默认=2048）。
        dropout: dropout的值（默认=0.1）。
        activation: 中间层的激活函数，可选relu或gelu（默认=relu）。
        layer_norm_eps: 层归一化组件中的eps值（默认=1e-5）。
        batch_first: 如果设为`True`，则输入和输出张量将按照（batch, seq, feature）的形式提供。
            默认：`False`。

    Examples::
        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
        >>> src = torch.rand(10, 32, 512)
        >>> out = encoder_layer(src)

    Alternatively, when ``batch_first`` is ``True``:
        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8, batch_first=True)
        >>> src = torch.rand(32, 10, 512)
        >>> out = encoder_layer(src)

    __init__:
        self_attn: 使用 MultiheadAttention 实现的自注意力机制。
        linear1, linear2: 用于前馈网络的线性层。
        dropout: 用于前馈网络中的Dropout操作。
        norm1, norm2: 层归一化层，用于自注意力和前馈网络之后。
        dropout1, dropout2: 用于残差连接后的Dropout操作。
        activation: 激活函数，根据传入的activation参数选择。
    """
    __constants__ = ['batch_first']

    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation="relu",
                 layer_norm_eps=1e-5, batch_first=False,
                 device=None, dtype=None) -> None:
        factory_kwargs = {'device': device, 'dtype': dtype}
        super(TransformerEncoderLayer, self).__init__()
        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=batch_first,
                                            **factory_kwargs)
        # Implementation of Feedforward model
        self.linear1 = Linear(d_model, dim_feedforward, **factory_kwargs)
        self.dropout = Dropout(dropout)
        self.linear2 = Linear(dim_feedforward, d_model, **factory_kwargs)

        self.norm1 = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
        self.norm2 = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
        self.dropout1 = Dropout(dropout)
        self.dropout2 = Dropout(dropout)

        self.activation = _get_activation_fn(activation)

    def __setstate__(self, state):
        if 'activation' not in state:
            state['activation'] = F.relu
        super(TransformerEncoderLayer, self).__setstate__(state)

    def forward(self, src: Tensor, src_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
        r"""Pass the input through the encoder layer.

        Args:
            src: the sequence to the encoder layer (required).
            src_mask: the mask for the src sequence (optional).
            src_key_padding_mask: the mask for the src keys per batch (optional).

        Shape:
            see the docs in Transformer class.

        forword flow:
            首先通过自注意力层(self_attn)处理输入src，得到src2。
            将src与经过Dropout操作的src2相加，然后通过层归一化(norm1)。
            接着，src通过前馈网络，依次经过linear1、激活函数、dropout、linear2、dropout2，再与src相加，最后通过层归一化(norm2)。
        """
        src2 = self.self_attn(src, src, src, attn_mask=src_mask,
                              key_padding_mask=src_key_padding_mask)[0]
        src = src + self.dropout1(src2)
        src = self.norm1(src)
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
        src = src + self.dropout2(src2)
        src = self.norm2(src)
        return src


class TransformerDecoderLayer(Module):
    r"""`TransformerDecoderLayer` 由自注意力(self-attn)、多头注意力(multi-head-attn)和前馈网络(feedforward network)组成。
    这个标准的解码器层基于论文 "Attention Is All You Need"。
    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
    Lukasz Kaiser, 和 Illia Polosukhin。2017年。注意力就是你所需要的一切。在《神经信息处理系统进展》期刊中，
    第6000至6010页。用户在应用过程中可以对其进行修改或以不同方式实现。

    参数:
        d_model: 输入中预期特征的数量（必需）。
        nhead: 多头注意力模型中的头数（必需）。
        dim_feedforward: 前馈网络模型的维度（默认=2048）。
        dropout: dropout的值（默认=0.1）。
        activation: 中间层的激活函数，可选relu或gelu（默认=relu）。
        layer_norm_eps: 层归一化组件中的eps值（默认=1e-5）。
        batch_first: 如果设为`True`，则输入和输出张量将按照（batch, seq, feature）的形式提供。
            默认：`False`。

    Examples::
        >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
        >>> memory = torch.rand(10, 32, 512)
        >>> tgt = torch.rand(20, 32, 512)
        >>> out = decoder_layer(tgt, memory)

    Alternatively, when ``batch_first`` is ``True``:
        >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8, batch_first=True)
        >>> memory = torch.rand(32, 10, 512)
        >>> tgt = torch.rand(32, 20, 512)
        >>> out = decoder_layer(tgt, memory)

    __init__:
        self_attn: 自注意力机制，用于处理目标序列内部的关系。
        multihead_attn: 多头注意力机制，用于处理目标序列和记忆序列之间的关系。
        linear1, linear2: 前馈网络的两个线性层。
        dropout: 用于前馈网络中的Dropout操作。
        norm1, norm2, norm3: 三个层归一化层，分别应用于自注意力、多头注意力和前馈网络之后。
        dropout1, dropout2, dropout3: 用于残差连接后的Dropout操作。
        activation: 激活函数，根据activation参数确定。
    """
    __constants__ = ['batch_first']

    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation="relu",
                 layer_norm_eps=1e-5, batch_first=False, device=None, dtype=None) -> None:
        factory_kwargs = {'device': device, 'dtype': dtype}
        super(TransformerDecoderLayer, self).__init__()
        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=batch_first,
                                            **factory_kwargs)
        self.multihead_attn = MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=batch_first,
                                                 **factory_kwargs)
        # Implementation of Feedforward model
        self.linear1 = Linear(d_model, dim_feedforward, **factory_kwargs)
        self.dropout = Dropout(dropout)
        self.linear2 = Linear(dim_feedforward, d_model, **factory_kwargs)

        self.norm1 = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
        self.norm2 = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
        self.norm3 = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
        self.dropout1 = Dropout(dropout)
        self.dropout2 = Dropout(dropout)
        self.dropout3 = Dropout(dropout)

        self.activation = _get_activation_fn(activation)

    def __setstate__(self, state):
        if 'activation' not in state:
            state['activation'] = F.relu
        super(TransformerDecoderLayer, self).__setstate__(state)

    def forward(self, tgt: Tensor, memory: Tensor, tgt_mask: Optional[Tensor] = None, memory_mask: Optional[Tensor] = None,
                tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:
        r"""使输入（及掩码）通过解码器层进行处理。

        参数:
            tgt: 需要送入解码器层的序列（必需）。
            memory: 来自编码器最后一层的序列（必需）。
            tgt_mask: 对于目标序列的掩码（可选）。
            memory_mask: 对于记忆序列的掩码（可选）。
            tgt_key_padding_mask: 每个批次的目标序列键的掩码（可选）。
            memory_key_padding_mask: 每个批次的记忆序列键的掩码（可选）。

        形状:
            请参阅Transformer类中的文档。

        流程：
            首先，目标序列通过自注意力机制处理，结果与原目标序列残差连接后进行层归一化。
            然后，处理后的目标序列与记忆序列通过多头注意力机制交互，结果与目标序列残差连接后进行层归一化。
            最后，目标序列通过前馈神经网络，包括线性层、激活函数、Dropout和另一个线性层，处理后的结果与目标序列残差连接后进行层归一化。
        """
        tgt2 = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask,
                              key_padding_mask=tgt_key_padding_mask)[0]
        tgt = tgt + self.dropout1(tgt2)
        tgt = self.norm1(tgt)
        tgt2 = self.multihead_attn(tgt, memory, memory, attn_mask=memory_mask,
                                   key_padding_mask=memory_key_padding_mask)[0]
        tgt = tgt + self.dropout2(tgt2)
        tgt = self.norm2(tgt)
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
        tgt = tgt + self.dropout3(tgt2)
        tgt = self.norm3(tgt)
        return tgt


def _get_clones(module, N):
    return ModuleList([copy.deepcopy(module) for i in range(N)])


def _get_activation_fn(activation):
    if activation == "relu":
        return F.relu
    elif activation == "gelu":
        return F.gelu

    raise RuntimeError("activation should be relu/gelu, not {}".format(activation))