解析Torch中多头注意力`MultiheadAttention`

前沿： 这部分内容是《Attention Is All You Need》出来之后就深入研究了这篇文章的模型结构，也是之后工作那一年进行实际落地的一小部分内容。最近再次使用它，顺带读了torch官方的实现，大家风范的实现，注意很多细节，值得我们学习，也顺带放在这，之后就不再了解这块内容了，过去式了。下面是内容:

MultiheadAttention

class MultiheadAttention(Module):
    r"""Allows the model to jointly attend to information
    from different representation subspaces.
    See `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_

    .. math::
        \text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O

    where :math:`head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)`.

    Args:
        embed_dim: total dimension of the model.
        num_heads: parallel attention heads.
        dropout: a Dropout layer on attn_output_weights. Default: 0.0.
        bias: add bias as module parameter. Default: True.
        add_bias_kv: add bias to the key and value sequences at dim=0.
        add_zero_attn: add a new batch of zeros to the key and
                       value sequences at dim=1.
        kdim: total number of features in key. Default: None.
        vdim: total number of features in value. Default: None.
        batch_first: If ``True``, then the input and output tensors are provided
            as (batch, seq, feature). Default: ``False`` (seq, batch, feature).

    Note that if :attr:`kdim` and :attr:`vdim` are None, they will be set
    to :attr:`embed_dim` such that query, key, and value have the same
    number of features.

    Examples::

        >>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
        >>> attn_output, attn_output_weights = multihead_attn(query, key, value)
    """
    __constants__ = ['batch_first']
    bias_k: Optional[torch.Tensor]
    bias_v: Optional[torch.Tensor]

    def __init__(self, embed_dim, num_heads, dropout=0., bias=True, add_bias_kv=False, add_zero_attn=False,
                 kdim=None, vdim=None, batch_first=False, device=None, dtype=None) -> None:
        """

        :param embed_dim: 模型的嵌入维度，即查询、键和值的维度
        :param num_heads: 并行的注意力头的数量
        :param dropout: 一个可选参数，用于定义在注意力权重上的 Dropout 层的丢弃率，默认为0
        :param bias: 一个布尔值，用于确定是否在投影权重中添加偏置，默认为 True
        :param add_bias_kv: 一个布尔值，用于确定是否在键和值的序列中添加额外的偏置项，默认为 False
        :param add_zero_attn: 一个布尔值，用于确定是否在键和值序列中添加一个新的零向量，默认为 False
        :param kdim: 键的维度，如果没有提供则默认等于 embed_dim
        :param vdim: 值的维度，如果没有提供则默认等于 embed_dim
        :param batch_first: 一个布尔值，用于确定输入和输出张量的维度顺序是否为 (batch, seq, feature)，默认为 False，即 (seq, batch, feature)
        :param device: 用于指定参数的设备（CPU/GPU）
        :param dtype: 用于指定参数的数据类型
        """
        # 创建一个字典 factory_kwargs 来存储设备和数据类型的信息，这些信息将在后面创建参数时使用。
        factory_kwargs = {
   
   'device': device, 'dtype': dtype}
        super(MultiheadAttention, self).__init__()

        #  设置 embed_dim、kdim 和 vdim 的值，如果 kdim 和 vdim 没有指定，则它们等于 embed_dim
        self.embed_dim = embed_dim
        self.kdim = kdim if kdim is not None else embed_dim
        self.vdim = vdim if vdim is not None else embed_dim

        # 判断键、值和嵌入维度是否相同，这影响到投影权重的创建方式
        self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim

        # 设置 num_heads、dropout 和 batch_first。
        # 计算每个头的维度 head_dim。
        # 确保 embed_dim 能够被 num_heads 整除，这是多头注意力机制的要求
        self.num_heads = num_heads
        self.dropout = dropout
        self.batch_first = batch_first
        self.head_dim = embed_dim // num_heads
        assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"

        # 根据 _qkv_same_embed_dim 的真假创建投影权重。如果维度相同，则创建一个大型的共享权重矩阵；如果不相同，则为查询、键和值创建单独的权重矩阵
        if self._qkv_same_embed_dim is False:
            self.q_proj_weight = Parameter(torch.empty((embed_dim, embed_dim), **factory_kwargs))
            self.k_proj_weight = Parameter(torch.empty((embed_dim, self.kdim), **factory_kwargs))
            self.v_proj_weight = Parameter(torch.empty((embed_dim, self.vdim), **factory_kwargs))
            self.register_parameter('in_proj_weight', None)
        else:
            self.in_proj_weight = Parameter(torch.empty((3 * embed_dim, embed_dim), **factory_kwargs))
            self.register_parameter('q_proj_weight', None)
            self.register_parameter('k_proj_weight', None)
            self.register_parameter('v_proj_weight', None)

        # 如果 bias 为 True，则创建一个偏置参数；否则，不创建偏置参数
        if bias:
            self.in_proj_bias = Parameter(torch.empty(3 * embed_dim, **factory_kwargs))
        else:
            self.register_parameter('in_proj_bias', None)

        # 创建输出投影层，这是一个线性层，用于在多头注意力计算之后将输出映射回 embed_dim 维度
        self.out_proj = NonDynamicallyQuantizableLinear(embed_dim, embed_dim, bias=bias, **factory_kwargs)

        # 如果 add_bias_kv 为 True，则创建键和值的偏置参数；否则，不创建这些偏置参数
        if add_bias_kv:
            self.bias_k = Parameter(torch.empty((1, 1, embed_dim), **factory_kwargs))
            self.bias_v = Parameter(torch.empty((1, 1, embed_dim), **factory_kwargs))
        else:
            self.bias_k = self.bias_v = None

        # 设置 add_zero_attn，用于控制是否在键和值序列中添加零向量
        self.add_zero_attn = add_zero_attn

        # 用 _reset_parameters 方法来初始化所有的参数，例如通过 Xavier 初始化等方法来确保参数的合理范围。这一步是确保模型在训练开始时有一个良好的初始状态
        self._reset_parameters()

    def _reset_parameters(self):
        if self._qkv_same_embed_dim:
            xavier_uniform_(self.in_proj_weight)
        else:
            xavier_uniform_(self.q_proj_weight)
            xavier_uniform_(self.k_proj_weight)
            xavier_uniform_(self.v_proj_weight)

        if self.in_proj_bias is not None:
            constant_(self.in_proj_bias, 0.)
            constant_(self.out_proj.bias, 0.)
        if self.bias_k is not None:
            xavier_normal_(self.bias_k)
        if self.bias_v is not None:
            xavier_normal_(self.bias_v)

    def __setstate__(self, state):
        # Support loading old MultiheadAttention checkpoints generated by v1.1.0
        if '_qkv_same_embed_dim' not in state:
            state['_qkv_same_embed_dim'] = True

        super(MultiheadAttention, self).__setstate__(state)

    def forward(self, query: Tensor, key: Tensor, value: Tensor, key_padding_mask: Optional[Tensor] = None,
                need_weights: bool = True, attn_mask: Optional[Tensor] = None) -> Tuple[Tensor, Optional[Tensor]]:
        r"""
        参数说明：
            query（查询）、key（键）、value（值）：将一个查询和一组键值对映射到一个输出。更多细节请参考论文《Attention Is All You Need》。
            key_padding_mask：如果提供，那么在键中的指定填充元素将被注意力机制忽略。当给定一个二进制掩码并且某个值为True时，注意力层上对应的值将被忽略。当给定一个字节掩码且某个值非零时，注意力层上对应的值也将被忽略。
            need_weights：输出注意力权重attn_output_weights。
            attn_mask：2D或3D的掩码，用于阻止对某些位置的注意力。一个2D掩码将为所有批次广播，而一个3D掩码允许为每个批次的条目指定不同的掩码。

        输入张量的形状：
            query：数学表示为(L, N, E)，其中L是目标序列长度，N是批次大小，E是嵌入维度。如果batch_first设置为True，则表示为(N, L, E)。
            key：数学表示为(S, N, E)