关于attention的学习（原理+代码）_nondynamicallyquantizablelinear

日志贴，有错拜托评论区指正~

附上李宏毅老师对transformer的讲解：

李宏毅2020机器学习深度学习(完整版)国语_哔哩哔哩_bilibili

关键的原理性问题

1、attention背景（没那么重要）：

seq2seq任务(如翻译、对话等)，输入一个序列得到一个序列，常用RNN实现，单向RNN可以看到序列当前位置及以前的内容、得到当前结果，双向RNN则可以看遍整个序列得到每一个结果。但是RNN不利于并行化，高度的并行化CNN又不能看遍序列，所以提出了self-attention layer.常见attention functions分为additive attention和dot-product attention，后者即softmax(QKT)V，便于矩阵并行化所以更优。

2、我之前不明所以以为transformer和attention是一个东西...实则不然。

 如图，是transformer的结构，由一个encoder（左）和一个decoder（右）组成，而encoder、decoder内部包含了attention layer，除了attention layer还有FFN（feed forward network，即线性+激活）

把encoder和decoder拆开用，多个encoder组合就是BERT模型，多个decoder组合就是GPT模型。

3、为什么transformer起源于NLP，在CV领域也用得很好？

嗯，这个是我项目答辩时候被老师问住的问题......我想说大白话，transformer在NLP很好用，于是就被拿来CV试试结果人家真的很好用...怕被老师骂，最终讷讷不得言。

Transformer在CV界火的原因是？_公众号机器学习与生成对抗网络的博客-CSDN博客

——也许这个问题的答案可以参考这个博主的文章。

在我导师看来，在当下的CV，attention层就应该和卷积层一样成为基础知识，被当作基本的层来使用。卷积更提取局部特征，attention则加强全局联系，扩大感受野。按照李宏毅老师的说法，CNN是简化版的self-attention，卷积只attend人工划定的感受野；self-attention先找出相关的像素，即自动学出感受野的形状和大小。

怎么将transformer用在CV里面呢？就是把图像也看作序列一样，将H x W x F的图像改为HW x F的维度，即可看作HW个F维向量。对于高像素图，通过打patch的操作将原图切作更小像素级的图像块送入attention层中进行计算。

4、Position Encoding

由于self-attention没有position information，所以加上通过学习得到的PE（维度和embedding的维度一样）来区分query的不同位置，或者说对相对位置进行表达。NLP中这个向量能决定当前词的位置，或者说在一个句子中不同的词之间的距离；CV中就是确定某个像素点的所在行列位置。

代码实现

（我用的是cuda10.2，torch1.10.0，python3.8.12，不同版本调用的代码可能不一样）

multihead_attn = torch.nn.MultiheadAttention(d_model, nhead, dropout)
out, attention_map = multihead_attn(query, key, value, attn_mask, key_padding_mask)

1、torch.nn.MultiheadAttention代码所在位置：torch\nn\modules\activation.py

（1）__init__函数中，必选输入参数有embed_dim和num_heads。

默认情况下q,k,v的embedding维度需要一样，即q,k的维度分别为(L, B, embed_dim)与(S, B, embed_dim)，否则应当输入特殊的参数kdim和vdim，即k维度为(S, B, kdim).

无论q,k,v的embedding维度一致与否，都会通过线性变换，变作同样的embed_dim维度，只是线性变换权重要分开放。

class MultiheadAttention(Module):
 
    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:
        factory_kwargs = {'device': device, 'dtype': dtype}
        super(MultiheadAttention, self).__init__()
        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
 
        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"
        # embed_dim必须能被num_heads整除
 
        # 以下proj_weight记作A
        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)
        # 以下in_proj_bias记作b
        if bias:
            self.in_proj_bias = Parameter(torch.empty(3 * embed_dim, **factory_kwargs))
        else:
            self.register_parameter('in_proj_bias', None)
        # forward过程第一步按公式y = xA^T + b对输入的q,k,v进行线性变换

（2）关于NonDynamicallyQuantizableLinear这个类，仅仅是为了避免在编写不当量化的注意力层脚本时触发一个不明显的错误。这里的out_proj是初始化了embed_dim到embed_dim的线性变化层的权重和偏差，用来对attention最后结果输出前做一次线性变换。

        self.out_proj = NonDynamicallyQuantizableLinear(embed_dim, embed_dim, bias=bias, **factory_kwargs)
 
        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
 
        self.add_zero_attn = add_zero_attn
 
        self._reset_parameters()
 
    # 让attention要训的参数好好初始化
    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)

（3）forward函数

关于参数和输出的注释。

    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"""
    Args:
        # 如果__init__时batch_first=True，则q、k、v的B在第一个维度，否则如下所示：
        query: (L, B, E_q)
        key: (S, B, E_k)
        value: (S, B, E_v)
        key_padding_mask: If specified, a mask of (B, S) indicating which elements within key to ignore for the purpose of attention. Binary and byte masks are supported.
            For a binary mask, a True value indicates that the corresponding key value will be ignored for the purpose of attention. For a byte mask, a non-zero value indicates that the corresponding key value will be ignored.
        need_weights: Default: True.
        attn_mask: If specified, a 2D or 3D mask preventing attention to certain positions. Must be of shape (L, S) or (B\cdot\text{num\_heads}, L, S). 
            A 2D mask will be broadcasted across the batch while a 3D mask allows for a different mask for each entry in the batch.
            Binary, byte, and float masks are supported. For a binary mask, a True value indicates that the corresponding position is not allowed to attend. For a byte mask, a non-zero value indicates that the corresponding position is not allowed to attend. For a float mask, the mask values will be added to the attention weight.
    Outputs:
        - **attn_output** - Attention outputs (L, B, E)或(B, L, E)
        - **attn_output_weights** - Attention output weights (B, L, S) when need_weights=True
        """

具体实现依靠F.multi_head_attention_forward

 
        if self.batch_first:
            query, key, value = [x.transpose(1, 0) for x in (query, key, value)]
        # 保证q,k,v的B一直在第二个维度
 
        
        # 将前面初始化好的权重偏差以及Q/K/V、mask传入更底层的函数，得到输出的最终结果和中间结果
        if not self._qkv_same_embed_dim:
            # 如果要用embedding维度不一样的k/v，需要令F.multi_head_attention_forward                        
            # 的输入参数use_separate_proj_weight=True
 
            attn_output, attn_output_weights = F.multi_head_attention_forward(
                query, key, value, self.embed_dim, self.num_heads,
                self.in_proj_weight, self.in_proj_bias,
                self.bias_k, self.bias_v, self.add_zero_attn,
                self.dropout, self.out_proj.weight, self.out_proj.bias,
                training=self.training,
                key_padding_mask=key_padding_mask, need_weights=need_weights,
                attn_mask=attn_mask, use_separate_proj_weight=True,
                q_proj_weight=self.q_proj_weight, k_proj_weight=self.k_proj_weight,
                v_proj_weight=self.v_proj_weight)
 
        else:
            attn_output, attn_output_weights = F.multi_head_attention_forward(
                query, key, value, self.embed_dim, self.num_heads,
                self.in_proj_weight, self.in_proj_bias,
                self.bias_k, self.bias_v, self.add_zero_attn,
                self.dropout, self.out_proj.weight, self.out_proj.bias,
                training=self.training,
                key_padding_mask=key_padding_mask, need_weights=need_weights,
                attn_mask=attn_mask)
        if self.batch_first:
            return attn_output.transpose(1, 0), attn_output_weights
        else:
            return attn_output, attn_output_weights

2、F.multi_head_attention_forward的实现代码位于torch\nn\functional.py中

（1）第一段无重点

def multi_head_attention_forward(
    query: Tensor,
    key: Tensor,
    value: Tensor,
    embed_dim_to_check: int,
    num_heads: int,
    in_proj_weight: Tensor,
    in_proj_bias: Optional[Tensor],
    bias_k: Optional[Tensor],
    bias_v: Optional[Tensor],
    add_zero_attn: bool,
    dropout_p: float,
    out_proj_weight: Tensor,
    out_proj_bias: Optional[Tensor],
    training: bool = True,
    key_padding_mask: Optional[Tensor] = None,
    need_weights: bool = True,
    attn_mask: Optional[Tensor] = None,
    use_separate_proj_weight: bool = False,
    q_proj_weight: Optional[Tensor] = None,
    k_proj_weight: Optional[Tensor] = None,
    v_proj_weight: Optional[Tensor] = None,
    static_k: Optional[Tensor] = None,
    static_v: Optional[Tensor] = None,
) -> Tuple[Tensor, Optional[Tensor]]:
    
    tens_ops = (query, key, value, in_proj_weight, in_proj_bias, bias_k, bias_v, out_proj_weight, out_proj_bias)
    # 这个看不懂，暂且跳过不影响
    if has_torch_function(tens_ops):
        return handle_torch_function(
            multi_head_attention_forward,
            tens_ops,
            query,
            key,
            value,
            embed_dim_to_check,
            num_heads,
            in_proj_weight,
            in_proj_bias,
            bias_k,
            bias_v,
            add_zero_attn,
            dropout_p,
            out_proj_weight,
            out_proj_bias,
            training=training,
            key_padding_mask=key_padding_mask,
            need_weights=need_weights,
            attn_mask=attn_mask,
            use_separate_proj_weight=use_separate_proj_weight,
            q_proj_weight=q_proj_weight,
            k_proj_weight=k_proj_weight,
            v_proj_weight=v_proj_weight,
            static_k=static_k,
            static_v=static_v,
        )
 
    # set up shape vars
    tgt_len, bsz, embed_dim = query.shape
    src_len, _, _ = key.shape
    assert embed_dim == embed_dim_to_check, \
        f"was expecting embedding dimension of {embed_dim_to_check}, but got {embed_dim}"
    if isinstance(embed_dim, torch.Tensor):
        # embed_dim can be a tensor when JIT tracing
        head_dim = embed_dim.div(num_heads, rounding_mode='trunc')
    else:
        head_dim = embed_dim // num_heads
    assert head_dim * num_heads == embed_dim, f"embed_dim {embed_dim} not divisible by num_heads {num_heads}"
    # 这里再次规定了embed_dim必需能被num_heads整除，否则会报错。 

（2）in_projection：将q,k,v都做一次线性变换，无论先前如何，变换后都是同样的emdding维度。

    # use_separate_proj_weight=True时，是不同embedding维度的kv输入
    if use_separate_proj_weight:
        # allow MHA to have different embedding dimensions when separate projection weights are used
        assert key.shape[:2] == value.shape[:2], \
            f"key's sequence and batch dims {key.shape[:2]} do not match value's {value.shape[:2]}"
        # kv的sequence和batch维度必须一致，即前两个维度必须都是(S,B)
    else:
        assert key.shape == value.shape, f"key shape {key.shape} does not match value shape {value.shape}"
 
    # 计算in-projection
    if not use_separate_proj_weight:
        q, k, v = _in_projection_packed(query, key, value, in_proj_weight, in_proj_bias)
    else:
        assert q_proj_weight is not None, "use_separate_proj_weight is True but q_proj_weight is None"
        assert k_proj_weight is not None, "use_separate_proj_weight is True but k_proj_weight is None"
        assert v_proj_weight is not None, "use_separate_proj_weight is True but v_proj_weight is None"
        if in_proj_bias is None:
            b_q = b_k = b_v = None
        else:
            b_q, b_k, b_v = in_proj_bias.chunk(3)
        q, k, v = _in_projection(query, key, value, q_proj_weight, k_proj_weight, v_proj_weight, b_q, b_k, b_v)

（2.1）_in_projection_packed和 _in_projection函数

见注释，前者是对embedding维度相同的q/k/v进行线性变换；后者是对维度不同的做。

def _in_projection_packed(
    q: Tensor,
    k: Tensor,
    v: Tensor,
    w: Tensor,
    b: Optional[Tensor] = None,
) -> List[Tensor]:
 
    E = q.size(-1)
    if k is v:
        if q is k:
        # q=k=v,做的是self-attention，可以直接将(3E,E)的权重矩阵w与(L,B,E)的q送入linear，        
        # linear的做法是q*w^T+b，所以得到结果(L,B,3E)。
        # 再用chunk在最后一维均分3块，得到3个(L,B,E)大小的q、k、v.
            return linear(q, w, b).chunk(3, dim=-1)
        else:
            # k=v, encoder-decoder attention,则k、v的linear变换可合并，q单独做
            w_q, w_kv = w.split([E, E * 2])
            if b is None:
                b_q = b_kv = None
            else:
                b_q, b_kv = b.split([E, E * 2])
            return (linear(q, w_q, b_q),) + linear(k, w_kv, b_kv).chunk(2, dim=-1)
    else:
        # q/k/v各不同，则先将w分3块，再分别做linear
        w_q, w_k, w_v = w.chunk(3)
        if b is None:
            b_q = b_k = b_v = None
        else:
            b_q, b_k, b_v = b.chunk(3)
        return linear(q, w_q, b_q), linear(k, w_k, b_k), linear(v, w_v, b_v)
 
 
def _in_projection(
    q: Tensor,
    k: Tensor,
    v: Tensor,
    w_q: Tensor,
    w_k: Tensor,
    w_v: Tensor,
    b_q: Optional[Tensor] = None,
    b_k: Optional[Tensor] = None,
    b_v: Optional[Tensor] = None,
) -> Tuple[Tensor, Tensor, Tensor]:
    
    # embedding维度上q,k,v不同，权重矩阵单独存入，检验输入输出大小是否符合后，分别做linear
    Eq, Ek, Ev = q.size(-1), k.size(-1), v.size(-1)
    assert w_q.shape == (Eq, Eq), f"expecting query weights shape of {(Eq, Eq)}, but got {w_q.shape}"
    assert w_k.shape == (Eq, Ek), f"expecting key weights shape of {(Eq, Ek)}, but got {w_k.shape}"
    assert w_v.shape == (Eq, Ev), f"expecting value weights shape of {(Eq, Ev)}, but got {w_v.shape}"
    assert b_q is None or b_q.shape == (Eq,), f"expecting query bias shape of {(Eq,)}, but got {b_q.shape}"
    assert b_k is None or b_k.shape == (Eq,), f"expecting key bias shape of {(Eq,)}, but got {b_k.shape}"
    assert b_v is None or b_v.shape == (Eq,), f"expecting value bias shape of {(Eq,)}, but got {b_v.shape}"
    return linear(q, w_q, b_q), linear(k, w_k, b_k), linear(v, w_v, b_v)

（3）检验mask输入符合数据类型和维度大小的要求否

    # prep attention mask
    if attn_mask is not None:
        if attn_mask.dtype == torch.uint8:
            warnings.warn("Byte tensor for attn_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")
            attn_mask = attn_mask.to(torch.bool)
        else:
            assert attn_mask.is_floating_point() or attn_mask.dtype == torch.bool, \
                f"Only float, byte, and bool types are supported for attn_mask, not {attn_mask.dtype}"
        # ensure attn_mask's dim is 3
        if attn_mask.dim() == 2:
            correct_2d_size = (tgt_len, src_len)
            if attn_mask.shape != correct_2d_size:
                raise RuntimeError(f"The shape of the 2D attn_mask is {attn_mask.shape}, but should be {correct_2d_size}.")
            attn_mask = attn_mask.unsqueeze(0)
        elif attn_mask.dim() == 3:
            correct_3d_size = (bsz * num_heads, tgt_len, src_len)
            if attn_mask.shape != correct_3d_size:
                raise RuntimeError(f"The shape of the 3D attn_mask is {attn_mask.shape}, but should be {correct_3d_size}.")
        else:
            raise RuntimeError(f"attn_mask's dimension {attn_mask.dim()} is not supported")
 
    # prep key padding mask
    if key_padding_mask is not None and key_padding_mask.dtype == torch.uint8:
        warnings.warn("Byte tensor for key_padding_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")
        key_padding_mask = key_padding_mask.to(torch.bool)

(4)k,v concat bias_k,bias_v

    # add bias along batch dimension (currently second)
    if bias_k is not None and bias_v is not None:
    # bias_k和bias_v在nn.MultiheadAttention中初始化为(1,1,E)大小的参数
        assert static_k is None, "bias cannot be added to static key."
        assert static_v is None, "bias cannot be added to static value."
        k = torch.cat([k, bias_k.repeat(1, bsz, 1)])  # （S+1, B ,E）
        v = torch.cat([v, bias_v.repeat(1, bsz, 1)])
        if attn_mask is not None:  # attention mask大小（L,S）或(?,L,S)
            attn_mask = pad(attn_mask, (0, 1))
        # pad操作是在mask最后一维上做padding，左侧一头不添，右侧一头添1。默认用0来pad。
        # 维度变为(L,S+1)
        if key_padding_mask is not None:
            key_padding_mask = pad(key_padding_mask, (0, 1))
    else:
        assert bias_k is None
        assert bias_v is None

（5）根据head数目改变q,k,v维度

    # reshape q, k, v for multihead attention and make em batch first
    q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
    # (B*H, L, E/H)
    if static_k is None:
        k = k.contiguous().view(k.shape[0], bsz * num_heads, head_dim).transpose(0, 1)
        #(B*H, S, E/H)或(B*H, S+1, E/H)
    else:
        # TODO finish disentangling control flow so we don't do in-projections when statics are passed
        assert static_k.size(0) == bsz * num_heads, \
            f"expecting static_k.size(0) of {bsz * num_heads}, but got {static_k.size(0)}"
        assert static_k.size(2) == head_dim, \
            f"expecting static_k.size(2) of {head_dim}, but got {static_k.size(2)}"
        k = static_k
    if static_v is None:
        v = v.contiguous().view(v.shape[0], bsz * num_heads, head_dim).transpose(0, 1)
        # 同k
    else:
        # TODO finish disentangling control flow so we don't do in-projections when statics are passed
        assert static_v.size(0) == bsz * num_heads, \
            f"expecting static_v.size(0) of {bsz * num_heads}, but got {static_v.size(0)}"
        assert static_v.size(2) == head_dim, \
            f"expecting static_v.size(2) of {head_dim}, but got {static_v.size(2)}"
        v = static_v

(6)add zero attention

对mask做一些处理

    # add zero attention along batch dimension (now first)
    if add_zero_attn:
        zero_attn_shape = (bsz * num_heads, 1, head_dim)
        k = torch.cat([k, torch.zeros(zero_attn_shape, dtype=k.dtype, device=k.device)], dim=1)
        # k(B*H, S, E/H)->(B*H, S+1, E/H) or ?
        v = torch.cat([v, torch.zeros(zero_attn_shape, dtype=v.dtype, device=v.device)], dim=1)
        if attn_mask is not None:
            attn_mask = pad(attn_mask, (0, 1))
        if key_padding_mask is not None:
            key_padding_mask = pad(key_padding_mask, (0, 1))
 
    # update source sequence length after adjustments
    src_len = k.size(1)  
    # S or S+1 or S+2? 默认情况下add_bias_kv=add_zero_attn=False，此处仍为S。
 
    # merge key padding and attention masks
    if key_padding_mask is not None:
        assert key_padding_mask.shape == (bsz, src_len), \
            f"expecting key_padding_mask shape of {(bsz, src_len)}, but got {key_padding_mask.shape}"
        key_padding_mask = key_padding_mask.view(bsz, 1, 1, src_len).   \
            expand(-1, num_heads, -1, -1).reshape(bsz * num_heads, 1, src_len)
        # key_mask(B*H, 1, S), attn_mask(B*H,L,S)
        if attn_mask is None:
            attn_mask = key_padding_mask
        elif attn_mask.dtype == torch.bool:
            attn_mask = attn_mask.logical_or(key_padding_mask)  
            # key_mask码掉的区域attn_mask也码掉
        else:
            attn_mask = attn_mask.masked_fill(key_padding_mask, float("-inf"))
 
    # convert mask to float
    if attn_mask is not None and attn_mask.dtype == torch.bool:
        new_attn_mask = torch.zeros_like(attn_mask, dtype=torch.float)
        new_attn_mask.masked_fill_(attn_mask, float("-inf"))
        attn_mask = new_attn_mask
 
    # adjust dropout probability，只有训练时设置dropout，推理时不用。
    if not training:
        dropout_p = 0.0

（7）用_scaled_dot_product_attention做attention

    attn_output, attn_output_weights = _scaled_dot_product_attention(q, k, v, attn_mask, dropout_p)
    attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
    # (L,B,E)
    attn_output = linear(attn_output, out_proj_weight, out_proj_bias)
    
    if need_weights:
        # average attention weights over heads
        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
        # attn_output_weights(B*H,L,S)->(B,H,L,S)->(B,L,S)
        return attn_output, attn_output_weights.sum(dim=1) / num_heads
    else:
        return attn_output, None

（7.1）_scaled_dot_product_attention

def _scaled_dot_product_attention(
    q: Tensor,
    k: Tensor,
    v: Tensor,
    attn_mask: Optional[Tensor] = None,
    dropout_p: float = 0.0,
) -> Tuple[Tensor, Tensor]:
  
    B, Nt, E = q.shape
    q = q / math.sqrt(E)
    # (B*H, L, E/H) x (B*H, E/H, S) -> attn(B*H, L, S)
    attn = torch.bmm(q, k.transpose(-2, -1))
    if attn_mask is not None:
        attn += attn_mask  # attn(B*H, L, S)
        # 在attention score上加attn_mask，mask的部分加负无穷大的数，经softmax后为0
    attn = softmax(attn, dim=-1)  # attn(B*H, L, S)
    # 在最后一个维度上做softmax
    if dropout_p > 0.0:
        attn = dropout(attn, p=dropout_p)
    # (B*H, L, S) x (B*H, S, E/H) -> (B*H, L, E/H)
    output = torch.bmm(attn, v)
    return output, attn