"未來"的經(jīng)典之作ViT：transformer is all you need！

點(diǎn)藍(lán)色字關(guān)注“機(jī)器學(xué)習(xí)算法工程師”

設(shè)為星標(biāo)，干貨直達(dá)！

從2020年，transformer開始在CV領(lǐng)域大放異彩：圖像分類（ViT, DeiT），目標(biāo)檢測（DETR，Deformable DETR），語義分割（SETR，MedT），圖像生成（GANsformer）等。而從深度學(xué)習(xí)暴發(fā)以來，CNN一直是CV領(lǐng)域的主流模型，而且取得了很好的效果，相比之下transformer卻獨(dú)霸NLP領(lǐng)域，transformer在CV領(lǐng)域的探索正是研究界想把transformer在NLP領(lǐng)域的成功借鑒到CV領(lǐng)域。對于圖像問題，CNN具有天然的先天優(yōu)勢（inductive bias）：平移不變性（translation equivariance）和局部性（locality）。而transformer雖然不并具備這些優(yōu)勢，但是transformer的核心self-attention的優(yōu)勢不像卷積那樣有固定且有限的感受野，self-attention操作可以獲得long-range信息（相比之下CNN要通過不斷堆積Conv layers來獲取更大的感受野），但訓(xùn)練的難度就比CNN要稍大一些。

ViT（vision transformer）是Google在2020年提出的直接將transformer應(yīng)用在圖像分類的模型，后面很多的工作都是基于ViT進(jìn)行改進(jìn)的。ViT的思路很簡單：直接把圖像分成固定大小的patchs，然后通過線性變換得到patch embedding，這就類比NLP的words和word embedding，由于transformer的輸入就是a sequence of token embeddings，所以將圖像的patch embeddings送入transformer后就能夠進(jìn)行特征提取從而分類了。ViT模型原理如下圖所示，其實(shí)ViT模型只是用了transformer的Encoder來提取特征（原始的transformer還有decoder部分，用于實(shí)現(xiàn)sequence to sequence，比如機(jī)器翻譯）。下面將分別對各個(gè)部分做詳細(xì)的介紹。

Patch Embedding

對于ViT來說，首先要將原始的2-D圖像轉(zhuǎn)換成一系列1-D的patch embeddings，這就好似NLP中的word embedding。輸入的2-D圖像記為，其中和分別是圖像的高和寬，而為通道數(shù)對于RGB圖像就是3。如果要將圖像分成大小為的patchs，可以通過reshape操作得到a sequence of patchs：，圖像共切分為個(gè)patchs，這也就是sequence的長度了，注意這里直接將patch拉平為1-D，其特征大小為。然后通過一個(gè)簡單的線性變換將patchs映射到大小的維度，這就是patch embeddings：，在實(shí)現(xiàn)上這等同于對進(jìn)行一個(gè)且stride為的卷積操作（雖然等同，但是ViT其實(shí)是不包含任何卷積操作的），下面是具體的實(shí)現(xiàn)代碼：

class PatchEmbed(nn.Module):
    """ Image to Patch Embedding
    """
    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = num_patches

        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        B, C, H, W = x.shape
        # FIXME look at relaxing size constraints
        assert H == self.img_size[0] and W == self.img_size[1], \
            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
        x = self.proj(x).flatten(2).transpose(1, 2)
        return x

Position Embedding

除了patch embeddings，模型還需要另外一個(gè)特殊的position embedding。transformer和CNN不同，需要position embedding來編碼tokens的位置信息，這主要是因?yàn)閟elf-attention是permutation-invariant，即打亂sequence里的tokens的順序并不會改變結(jié)果。如果不給模型提供patch的位置信息，那么模型就需要通過patchs的語義來學(xué)習(xí)拼圖，這就額外增加了學(xué)習(xí)成本。ViT論文中對比了幾種不同的position embedding方案(如下），最后發(fā)現(xiàn)如果不提供positional embedding效果會差，但其它各種類型的positional embedding效果都接近，這主要是因?yàn)閂iT的輸入是相對較大的patchs而不是pixels，所以學(xué)習(xí)位置信息相對容易很多。

• 無positional embedding
• 1-D positional embedding：把2-D的patchs看成1-D序列
• 2-D positional embedding：考慮patchs的2-D位置（x, y）
• Relative positional embeddings：patchs的相對位置

transformer原論文中是默認(rèn)采用固定的positional embedding，但ViT中默認(rèn)采用學(xué)習(xí)（訓(xùn)練的）的1-D positional embedding，在輸入transformer的encoder之前直接將patch embeddings和positional embedding相加:

# 這里多1是為了后面要說的class token，embed_dim即patch embed_dim
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim)) 

# patch emded + pos_embed
x = x + self.pos_embed

論文中也對學(xué)習(xí)到的positional embedding進(jìn)行了可視化，發(fā)現(xiàn)相近的patchs的positional embedding比較相似，而且同行或同列的positional embedding也相近：

這里額外要注意的一點(diǎn)，如果改變圖像的輸入大小，ViT不會改變patchs的大小，那么patchs的數(shù)量會發(fā)生變化，那么之前學(xué)習(xí)的pos_embed就維度對不上了，ViT采用的方案是通過插值來解決這個(gè)問題：

def resize_pos_embed(posemb, posemb_new):
    # Rescale the grid of position embeddings when loading from state_dict. Adapted from
    # https://github.com/google-research/vision_transformer/blob/00883dd691c63a6830751563748663526e811cee/vit_jax/checkpoint.py#L224
    _logger.info('Resized position embedding: %s to %s', posemb.shape, posemb_new.shape)
    ntok_new = posemb_new.shape[1]
    # 除去class token的pos_embed
    posemb_tok, posemb_grid = posemb[:, :1], posemb[0, 1:]
    ntok_new -= 1
    gs_old = int(math.sqrt(len(posemb_grid)))
    gs_new = int(math.sqrt(ntok_new))
    _logger.info('Position embedding grid-size from %s to %s', gs_old, gs_new)
    # 把pos_embed變換到2-D維度再進(jìn)行插值
    posemb_grid = posemb_grid.reshape(1, gs_old, gs_old, -1).permute(0, 3, 1, 2)
    posemb_grid = F.interpolate(posemb_grid, size=(gs_new, gs_new), mode='bilinear')
    posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(1, gs_new * gs_new, -1)
    posemb = torch.cat([posemb_tok, posemb_grid], dim=1)
    return posemb

但是這種情形一般會造成性能少許損失，可以通過finetune模型來解決。另外最新的論文CPVT通過implicit Conditional Position encoding來解決這個(gè)問題（插入Conv來隱式編碼位置信息，zero padding讓Conv學(xué)習(xí)到絕對位置信息）。

Class Token

除了patch tokens，ViT借鑒BERT還增加了一個(gè)特殊的class token。后面會說，transformer的encoder輸入是a sequence patch embeddings，輸出也是同樣長度的a sequence patch features，但圖像分類最后需要獲取image feature，簡單的策略是采用pooling，比如求patch features的平均來獲取image feature，但是ViT并沒有采用類似的pooling策略，而是直接增加一個(gè)特殊的class token，其最后輸出的特征加一個(gè)linear classifier就可以實(shí)現(xiàn)對圖像的分類（ViT的pre-training時(shí)是接一個(gè)MLP head），所以輸入ViT的sequence長度是。class token對應(yīng)的embedding在訓(xùn)練時(shí)隨機(jī)初始化，然后通過訓(xùn)練得到，具體實(shí)現(xiàn)如下：

# 隨機(jī)初始化
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))

# Classifier head
self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()

# 具體forward過程
B = x.shape[0]
x = self.patch_embed(x)
cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
x = torch.cat((cls_tokens, x), dim=1)
x = x + self.pos_embed

Transformer Encoder

transformer最核心的操作就是self-attention，其實(shí)attention機(jī)制很早就在NLP和CV領(lǐng)域應(yīng)用了，比如帶有attention機(jī)制的seq2seq模型，但是transformer完全摒棄RNN或LSTM結(jié)構(gòu)，直接采用attention機(jī)制反而取得了更好的效果：attention is all you need！簡單來說，attention就是根據(jù)當(dāng)前查詢對輸入信息賦予不同的權(quán)重來聚合信息，從操作上看就是一種“加權(quán)平均”。attention中共有3個(gè)概念：query, key和value，其中key和value是成對的，對于一個(gè)給定的query向量，通過內(nèi)積計(jì)算來匹配k個(gè)key向量（維度也是d，堆積起來即矩陣），得到的內(nèi)積通過softmax來歸一化得到k個(gè)權(quán)重，那么對于query其attention的輸出就是k個(gè)key向量對應(yīng)的value向量（即矩陣）的加權(quán)平均值。對于一系列的N個(gè)query（即矩陣），可以通過矩陣計(jì)算它們的attention輸出：

這里的為縮放因子以避免點(diǎn)積帶來的方差影響。上述的Attention機(jī)制稱為Scaled dot product attention，其實(shí)attention機(jī)制的變種有很多，但基本原理是相似的。如果都是從一個(gè)包含個(gè)向量的sequence（）通過線性變換得到：那么此時(shí)就變成了self-attention，這個(gè)時(shí)候就有個(gè)（key,value）對，那么。self-attention是transformer最核心部分，self-attention其實(shí)就是輸入向量之間進(jìn)行相互attention來學(xué)習(xí)到新特征。前面說過我們已經(jīng)得到圖像的patch sequence，那么送入self-attention就能到同樣size的sequence輸出，只不過特征改變了。

更進(jìn)一步，transformer采用的是multi-head self-attention (MSA），所謂的MSA就是采用定義h個(gè)attention heads，即采用h個(gè)self-attention應(yīng)用在輸入sequence上，在操作上可以將sequence拆分成h個(gè)size為的sequences，這里，h個(gè)不同的heads得到的輸出concat在一起然后通過線性變換得到最終的輸出，size也是：

MSA的計(jì)算量是和成正相關(guān)的，所以ViT的輸入是patch embeddings，而不是pixel embeddings，這有計(jì)算量上的考慮。在實(shí)現(xiàn)上，MSA是可以并行計(jì)算各個(gè)head的，具體代碼如下：

class Attention(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
    
        self.scale = qk_scale or 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)
        # 這里包含了dropout
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x):
        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)
        q, k, v = qkv[0], qkv[1], qkv[2]   # make torchscript happy (cannot use tensor as tuple)

        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x

在transformer中，MSA后跟一個(gè)FFN（Feed-forward network），這個(gè)FFN包含兩個(gè)FC層，第一個(gè)FC層將特征從維度變換成，后一個(gè)FC層將特征從維度恢復(fù)成，中間的非線性激活函數(shù)采用GeLU，其實(shí)這就是一個(gè)MLP，具體實(shí)現(xiàn)如下：

class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x

那么一個(gè)完成transformer encoder block就包含一個(gè)MSA后面接一個(gè)FFN，其實(shí)MSA和FFN均包含和ResNet一樣的skip connection，另外MSA和FFN后面都包含layer norm層，具體實(shí)現(xiàn)如下：

class Block(nn.Module):

    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = Attention(
            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)

    def forward(self, x):
        x = x + self.drop_path(self.attn(self.norm1(x)))
        x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x

ViT

對于ViT模型來說，就類似CNN那樣，不斷堆積transformer encoder blocks，最后提取class token對應(yīng)的特征用于圖像分類，論文中也給出了模型的公式表達(dá)，其中（1）就是提取圖像的patch embeddings，然后和class token對應(yīng)的embedding拼接在一起并加上positional embedding；（2）是MSA，而（3）是MLP，（2）和（3）共同組成了一個(gè)transformer encoder block，共有層；（4）是對class token對應(yīng)的輸出做layer norm，然后就可以用來圖像分類。

除了完全無卷積的ViT模型外，論文中也給出了Hybrid Architecture，簡單來說就是先用CNN對圖像提取特征，從CNN提取的特征圖中提取patch embeddings，CNN已經(jīng)將圖像降采樣了，所以patch size可以為。

ViT模型的超參數(shù)主要包括以下，這些超參數(shù)直接影響模型參數(shù)以及計(jì)算量：

1. Layers：block的數(shù)量；
2. Hidden size D：隱含層特征，D在各個(gè)block是一直不變的；
3. MLP size：一般設(shè)置為4D大??；
4. Heads：MSA中的heads數(shù)量；
5. Patch size：模型輸入的patch size，ViT中共有兩個(gè)設(shè)置：14x14和16x16，這個(gè)只影響計(jì)算量；

類似BERT，ViT共定義了3種不同大小的模型：Base，Large和Huge，其對應(yīng)的模型參數(shù)不同，如下所示。如ViT-L/16指的是采用Large結(jié)構(gòu)，輸入的patch size為16x16。

模型效果

ViT并不像CNN那樣具有inductive bias，論文中發(fā)現(xiàn)如果如果直接在ImageNet上訓(xùn)練，同level的ViT模型效果要差于ResNet，但是如果在比較大的數(shù)據(jù)集上petraining，然后再finetune，效果可以超越ResNet。比如ViT在Google私有的300M JFT數(shù)據(jù)集上pretrain后，在ImageNet上的最好Top-1 acc可達(dá)88.55%，這已經(jīng)和ImageNet上的SOTA相當(dāng)了（Noisy Student EfficientNet-L2效果為88.5%，Google最新的SOTA是Meta Pseudo Labels，效果可達(dá)90.2%）：

那么ViT至少需要多大的數(shù)據(jù)量才能和CNN旗鼓相當(dāng)呢？這個(gè)論文也做了實(shí)驗(yàn)，結(jié)果如下圖所示，從圖上所示這個(gè)預(yù)訓(xùn)練所使用的數(shù)據(jù)量要達(dá)到100M時(shí)才能顯示ViT的優(yōu)勢。transformer的一個(gè)特色是它的scalability：當(dāng)模型和數(shù)據(jù)量提升時(shí)，性能持續(xù)提升。在大數(shù)據(jù)面前，ViT可能會發(fā)揮更大的優(yōu)勢。

此外，論文中也對ViT做了進(jìn)一步分析，如分析了不同layers的mean attention distance，這個(gè)類比于CNN的感受野。論文中發(fā)現(xiàn)前面層的“感受野”雖然差異很大，但是總體相比后面層“感受野”較小，而模型后半部分“感受野”基本覆蓋全局，和CNN比較類似，說明ViT也最后學(xué)習(xí)到了類似的范式。

當(dāng)然，ViT還可以根據(jù)attention map來可視化模型具體關(guān)注圖像的哪個(gè)部分，從結(jié)果上看比較合理：

我個(gè)人覺得ViT算是一個(gè)很好的開始，雖然ViT也有一些問題，但是至少證明了純粹的transformer在CV領(lǐng)域應(yīng)用的可能性。近期也有一些后續(xù)的改進(jìn)工作，感興趣的可以進(jìn)一步了解：

• [DeiT] Training data-efficient image transformers & distillation through attention
• [T2T-ViT] Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet
• [CPVT] Do We Really Need Explicit Position Encodings for Vision Transformers?
• [PVT] Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions
• [TNT] Transformer in Transformer

參考

1. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
2. timm
3. Awesome-Visual-Transformer
4. Do We Really Need Explicit Position Encodings for Vision Transformers?