Vit的学习笔记

Posted 2021-12-22 songyuc

1 致谢

感谢paddle开源的Vit系列教程《飞桨——从零开始学视觉Transformer》；

2 Vit模型介绍

2.2 Attention模块（Multi-head Attention）

Attention模块是一种seq2seq的变换模块，其输入输出的特征维数相同；
参数：

• num_heads: attention head的数量，默认为8；将向量分成8份的话，计算效率也比较高。
• qkv_bias: 在qkv映射时使用偏置参数，默认为False，但是qkv_bias默认需要开启。
• attn_drop: self.attn_drop的丢弃率，默认为0。
• proj_drop: self.proj_drop的丢弃率，默认为0。

Note:

• qkv_bias在一般情况下必须开启，这样才符合逻辑回归的数学原理，我们在timm的Vit实现中也可以看到此参数默认是开启的，这里我们引用timm中Vit模型的参数解释来说明：

  qkv_bias: bool, If True, enable qkv(nn.Linear) layer with bias, default: True

  可知qkv_bias默认需要开启。

# timm库中的Attention实现
# 在torch中声明模块需要继承nn.Module
class Attention(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.):
        super().__init__()
        self.num_heads = num_heads
        # dim需要是num_heads的整数倍，可以参考HuggingFace的实现，
        # https://huggingface.co/transformers/_modules/transformers/models/vit/modeling_vit.html#ViTForImageClassification
        # TODO: assert dim % num_heads == 0
        head_dim = dim // num_heads
        self.scale = head_dim ** -0.5

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x):
    	# B为batch-size，
    	# N为patch-num，
    	# C为通道数，
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        # self.qkv() --> x: (B, N, 3C)
        # complex: 3C^2
        # reshape() --> x1: (B, N, 3, nhead, dhead)
        # complex: 0
        # permute() --> qkv: (3, B, nhead, N, dhead)
        # complex: 0
        q, k, v = qkv.unbind(0)   # make torchscript happy (cannot use tensor as tuple)

		# q: (B, nhead, N, dhead)
		# k: (B, nhead, N, dhead)
        attn = (q @ k.transpose(-2, -1)) * self.scale
        # q@k' --> attn: (B, nhead, N, N)
		# complex: B*nhead*N*dhead*N = BN^2*C
		# *self.scale --> attn
		# complex: B*nhead*N*N = BN^2*nhead
        attn = attn.softmax(dim=-1)
        # complex: BN^2*nhead
        attn = self.attn_drop(attn)
        # complex: 0
		
		# attn: (B, nhead, N, N)
		# v: (B, nhead, N, dhead)
        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        # attn @ v --> x
        # complex: B*nhead*N*N*dhead = BN^2C
        x = self.proj(x)
        x = self.proj_drop(x)
        return x

复杂度分析

可以看到，Attention模块的计算复杂度为
$$3C^2+B{N}^{2}C+B{N}^2\cdot nheads+B{N}^2$$
我们将复杂度按照算子分开列出：

• FC算子self.qkv: $3C^2$
• 乘法注意力q@k': $B{N}^{2}C$
• Softmax算子: $B{N}^2\cdot nheads$
• 注意力加权运算: $B{N}^2$

2.2 Vit模型为什么要使用Multi-head-Attention?

这里我们针对仅原始Vit的Multi-head-Attention进行讨论，并没有扩展到所有SOTA的vit类型算法；
由Attention模块的复杂度分析可知，仅有Softmax算子的复杂度$B{N}^2\cdot nheads$与注意力head的数量有关；于是我们可以知道使用Multi-head-Attention会使得Self-Attention算法的复杂度增加，且其增加的复杂度主要是来源于softmax算子的运算量增加了，且其运算量增加了$nheads$倍；