<input id="0qass"><u id="0qass"></u></input>
  • <input id="0qass"><u id="0qass"></u></input>
  • <menu id="0qass"><u id="0qass"></u></menu>

    PVT:可用于密集任務backbone的金字塔視覺transformer!

    機器學習系列 專欄收錄該內容
    13 篇文章 0 訂閱

    基于detectron2實現的PVT開源了,歡迎star:https://github.com/xiaohu2015/pvt_detectron2

    自從ViT之后,關于vision transformer的研究呈井噴式爆發,從思路上分主要沿著兩大個方向,一是提升ViT在圖像分類的效果;二就是將ViT應用在其它圖像任務中,比如分割和檢測任務上,這里介紹的PVT(Pyramid Vision Transformer) 就屬于后者。PVT相比ViT引入了和CNN類似的金字塔結構,使得PVT像CNN那樣作為backbone應用在dense prediction任務(分割和檢測等)。

    圖片

    CNN結構常用的是一種金字塔架構,如上圖所示,CNN網絡一般可以劃分為不同的stage,在每個stage開始時,特征圖的長和寬均減半,而特征維度(channel)擴寬2倍。這主要有兩個方面的考慮,一是采用stride=2的卷積或者池化層對特征降維可以增大感受野,另外也可以減少計算量,但同時空間上的損失用channel維度的增加來彌補。但是ViT本身就是全局感受野,所以ViT就比較簡單直接了,直接將輸入圖像tokens化后就不斷堆積相同的transformer encoders,這應用在圖像分類上是沒有太大的問題。但是如果應用在密集任務上,會遇到問題:一是分割和檢測往往需要較大的分辨率輸入,當輸入圖像增大時,ViT的計算量會急劇上升;二是ViT直接采用較大patchs進行token化,如采用16x16大小那么得到的粗粒度特征,對密集任務來說損失較大。這正是PVT想要解決的問題,PVT采用和CNN類似的架構,將網絡分成不同的stages,每個stage相比之前的stage特征圖的維度是減半的,這意味著tokens數量減少4倍,具體結構如下:

    圖片

    每個stage的輸入都是一個維度的3-D特征圖,對于第1個stage,輸入就是RGB圖像,對于其它stage可以將tokens重新reshape成3-D特征圖。在每個stage開始,首先像ViT一樣對輸入圖像進行token化,即進行patch embedding,patch大小均采用2x2大小(第1個stage的patch大小是4x4),這意味著該stage最終得到的特征圖維度是減半的,tokens數量對應減少4倍。PVT共4個stage,這和ResNet類似,4個stage得到的特征圖相比原圖大小分別是1/4,1/8,1/16和1/32。由于不同的stage的tokens數量不一樣,所以每個stage采用不同的position embeddings,在patch embed之后加上各自的position embedding,當輸入圖像大小變化時,position embeddings也可以通過插值來自適應。

    不同的stage的tokens數量不同,越靠前的stage的patchs數量越多,我們知道self-attention的計算量與sequence的長度的平方成正比,如果PVT和ViT一樣,所有的transformer encoders均采用相同的參數,那么計算量肯定是無法承受的。PVT為了減少計算量,不同的stages采用的網絡參數是不同的。PVT不同系列的網絡參數設置如下所示,這里為patch的size,為特征維度大小,為MHA(multi-head attention)的heads數量,為FFN的擴展系數,transformer中默認為4。

    圖片

    可以見到隨著stage,特征的維度是逐漸增加的,比如stage1的特征維度只有64,而stage4的特征維度為512,這種設置和常規的CNN網絡設置是類似的,所以前面stage的patchs數量雖然大,但是特征維度小,所以計算量也不是太大。不同體量的PVT其差異主要體現在各個stage的transformer encoder的數量差異。

    PVT為了進一步減少計算量,將常規的multi-head attention (MHA)用spatial-reduction attention (SRA)來替換。SRA的核心是減少attention層的key和value對的數量,常規的MHA在attention層計算時key和value對的數量為sequence的長度,但是SRA將其降低為原來的。SRA的具體結構如下所示:

    圖片

    在實現上,首先將維度為的patch embeddings通過 reshape變換到維度為的3-D特征圖,然后均分大小為的patchs,每個patchs通過線性變換將得到維度為的patch embeddings(這里實現上其實和patch emb操作類似,等價于一個卷積操作),最后應用一個layer norm層,這樣就可以大大降低K和V的數量。具體實現代碼如下:

    class Attention(nn.Module):
        def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., sr_ratio=1):
            super().__init__()
            assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}."
    
            self.dim = dim
            self.num_heads = num_heads
            head_dim = dim // num_heads
            self.scale = qk_scale or head_dim ** -0.5
    
            self.q = nn.Linear(dim, dim, bias=qkv_bias)
            self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias)
            self.attn_drop = nn.Dropout(attn_drop)
            self.proj = nn.Linear(dim, dim)
            self.proj_drop = nn.Dropout(proj_drop)
    
            self.sr_ratio = sr_ratio
            # 實現上這里等價于一個卷積層
            if sr_ratio > 1:
                self.sr = nn.Conv2d(dim, dim, kernel_size=sr_ratio, stride=sr_ratio)
                self.norm = nn.LayerNorm(dim)
    
        def forward(self, x, H, W):
            B, N, C = x.shape
            q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
    
            if self.sr_ratio > 1:
                x_ = x.permute(0, 2, 1).reshape(B, C, H, W)
                x_ = self.sr(x_).reshape(B, C, -1).permute(0, 2, 1) # 這里x_.shape = (B, N/R^2, C)
                x_ = self.norm(x_)
                kv = self.kv(x_).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
            else:
                kv = self.kv(x).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
            k, v = kv[0], kv[1]
    
            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
    
    

    從PVT的網絡設置上,前面的stage的取較大的值,比如stage1的,說明這里直接將Q和V的數量直接減為原來的1/64,這個就大大降低計算量了。

    PVT具體到圖像分類任務上,和ViT一樣也通過引入一個class token來實現最后的分類,不過PVT是在最后的一個stage才引入:

        def forward_features(self, x):
            B = x.shape[0]
    
            # stage 1
            x, (H, W) = self.patch_embed1(x)
            x = x + self.pos_embed1
            x = self.pos_drop1(x)
            for blk in self.block1:
                x = blk(x, H, W)
            x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
    
            # stage 2
            x, (H, W) = self.patch_embed2(x)
            x = x + self.pos_embed2
            x = self.pos_drop2(x)
            for blk in self.block2:
                x = blk(x, H, W)
            x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
    
            # stage 3
            x, (H, W) = self.patch_embed3(x)
            x = x + self.pos_embed3
            x = self.pos_drop3(x)
            for blk in self.block3:
                x = blk(x, H, W)
            x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
    
            # stage 4
            x, (H, W) = self.patch_embed4(x)
            cls_tokens = self.cls_token.expand(B, -1, -1) # 引入class token
            x = torch.cat((cls_tokens, x), dim=1)
            x = x + self.pos_embed4
            x = self.pos_drop4(x)
            for blk in self.block4:
                x = blk(x, H, W)
    
            x = self.norm(x)
    
            return x[:, 0]
    
    

    具體到分類任務上,PVT在ImageNet上的Top-1 Acc其實是和ViT差不多的。其實PVT最重要的應用是作為dense任務如分割和檢測的backbone,一方面PVT通過一些巧妙的設計使得對于分辨率較大的輸入圖像,其模型計算量不像ViT那么大,論文中比較了ViT-Small/16 ,ViT-Small,PVT-Small和ResNet50四種網絡在不同的輸入scale下的GFLOPs,可以看到PVT相比ViT要好不少,當輸入scale=640時,PVT-Small和ResNet50的計算量是類似的,但是如果到更大的scale,PVT的增長速度就遠超過ResNet50了。

    圖片

    PVT的另外一個相比ViT的優勢就是其可以輸出不同scale的特征圖,這對于分割和檢測都是非常重要的。因為目前大部分的分割和檢測模型都是采用FPN結構,而PVT這個特性可以使其作為替代CNN的backbone而無縫對接分割和檢測的heads。論文中做了大量的關于檢測,語義分割以及實例分割的實驗,可以看到PVT在dense任務的優勢。比如,在更少的推理時間內,基于PVT-Small的RetinaNet比基于R50的RetinaNet在COCO上的AP值更高(38.7 vs. 36.3),雖然繼續增加scale可以提升效果,但是就需要額外的推理時間:

    圖片

    所以雖然PVT可以解決一部分問題,但是如果輸入圖像分辨率特別大,可能基于CNN的方案還是最優的。另外曠視最新的一篇論文YOLOF指出其實ResNet一個C5特征加上一些增大感受野的模塊就可以在檢測上實現類似的效果,這不得不讓人思考多尺度特征是不是必須的,而且transformer encoder本身就是全局感受野的。近期Intel提出的DPT直接在ViT模型的基礎上通過Reassembles operation來得到不同scale的特征圖以用于dense任務,并在ADE20K語義分割數據集上達到新的SOTA(mIoU 49.02)。而在近日,微軟提出的Swin Transformer和PVT的網絡架構和很類似,但其性能在各個檢測和分割數據集上效果達到SOTA(在ADE20K語義分割數據集mIoU 53.5),其核心提出了一種shifted window方法來減少self-attention的計算量。

    相信未來會有更好的work!期待!

    參考

    1. Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions

    2. whai362/PVT

    3. 大白話Pyramid Vision Transformer

    4. You Only Look One-level Feature

    5. Swin Transformer: Hierarchical Vision Transformer using Shifted Windows

    6. Vision Transformers for Dense Prediction

    推薦閱讀

    谷歌提出Meta Pseudo Labels,刷新ImageNet上的SOTA!

    "未來"的經典之作ViT:transformer is all you need!

    大道至簡!深度解讀CVPR2021論文RepVGG!

    PyTorch 源碼解讀之 torch.autograd

    漲點神器FixRes:兩次超越ImageNet數據集上的SOTA

    Transformer為何能闖入CV界秒殺CNN?

    SWA:讓你的目標檢測模型無痛漲點1% AP

    CondInst:性能和速度均超越Mask RCNN的實例分割模型

    centerX: 用新的視角的方式打開CenterNet

    mmdetection最小復刻版(十一):概率Anchor分配機制PAA深入分析

    MMDetection新版本V2.7發布,支持DETR,還有YOLOV4在路上!

    CNN:我不是你想的那樣

    TF Object Detection 終于支持TF2了!

    無需tricks,知識蒸餾提升ResNet50在ImageNet上準確度至80%+

    不妨試試MoCo,來替換ImageNet上pretrain模型!

    重磅!一文深入深度學習模型壓縮和加速

    從源碼學習Transformer!

    mmdetection最小復刻版(七):anchor-base和anchor-free差異分析

    mmdetection最小復刻版(四):獨家yolo轉化內幕

    機器學習算法工程師


    一個用心的公眾號

    圖片


    圖片

    • 0
      點贊
    • 0
      評論
    • 0
      收藏
    • 一鍵三連
      一鍵三連
    • 掃一掃,分享海報

    打賞
    文章很值,打賞犒勞作者一下
    相關推薦
    ??2020 CSDN 皮膚主題: 編程工作室 設計師:CSDN官方博客 返回首頁

    打賞

    機器學習算法工程師

    你的鼓勵將是我創作的最大動力

    ¥2 ¥4 ¥6 ¥10 ¥20
    輸入1-500的整數
    余額支付 (余額:-- )
    掃碼支付
    掃碼支付:¥2
    獲取中
    掃碼支付

    您的余額不足,請更換掃碼支付或充值

    打賞作者

    實付
    使用余額支付
    點擊重新獲取
    掃碼支付
    錢包余額 0

    抵扣說明:

    1.余額是錢包充值的虛擬貨幣,按照1:1的比例進行支付金額的抵扣。
    2.余額無法直接購買下載,可以購買VIP、C幣套餐、付費專欄及課程。

    余額充值
    多乐彩