卷積視覺Transformer (CvT)

在本節中，我們將深入探討卷積視覺Transformer (CvT)，它是視覺Transformer (ViT)[1]的一種變體，並廣泛用於計算機視覺中的影像分類任務。

回顧

在深入CvT之前，讓我們先回顧一下前面章節中介紹的ViT架構，以便更好地理解CvT架構。ViT將每張影像分解為固定長度的token序列（即非重疊的塊），然後應用多個標準Transformer層，這些層由多頭自注意力和逐位置前饋模組（FFN）組成，以對分類的全域性關係進行建模。

概述

卷積視覺Transformer (CvT) 模型由 Haiping Wu、Bin Xiao、Noel Codella、Mengchen Liu、Xiyang Dai、Lu Yuan 和 Lei Zhang 在《CvT: Introducing Convolutions to Vision Transformers》[2]中提出。CvT 融合了 CNN 的所有優點：區域性感受野、權重共享和空間下采樣，以及平移、縮放、失真不變性，同時保留了 Transformer 的優點：動態注意力、全域性上下文融合和更好的泛化能力。與 ViT 相比，CvT 在保持計算效率的同時實現了卓越的效能。此外，由於卷積引入了內建的區域性上下文結構，CvT 不再需要位置編碼，這使其在適應需要可變輸入解析度的各種視覺任務方面具有潛在優勢。

架構

CvT Architecture (a) 整體架構，顯示了卷積Token嵌入層實現的層次化多階段結構。(b) 卷積Transformer塊的細節，其第一層包含卷積投影。[2]

上圖中的CvT架構圖展示了三階段管道的主要步驟。其核心是CvT將兩個基於卷積的操作融合到視覺Transformer架構中

卷積Token嵌入：想象一下將輸入影像分割成重疊的影像塊，將它們重塑為token，然後將它們輸入到卷積層。這減少了token的數量（類似於下采樣影像中的畫素），同時提高了它們的特徵豐富度，類似於傳統的CNN。與其他Transformer不同，我們跳過了向token新增預定義位置資訊，完全依賴卷積操作來捕獲空間關係。

Projection Layer (a) ViT 中的線性投影。(b) 卷積投影。(c) 壓縮卷積投影 (CvT 中的預設設定)。[2]

卷積Transformer塊：CvT 中的每個階段都包含一系列這樣的塊。在這裡，我們不使用 ViT 中常見的線性投影，而是使用深度可分離卷積（卷積投影）來處理自注意力模組的“查詢”、“鍵”和“值”元件，如上圖所示。這既保留了 Transformer 的優點，又提高了效率。請注意，“分類標記”（用於最終預測）僅在最後一個階段新增。最後，一個標準的完全連線層分析最終的分類標記以預測影像類別。

CvT 架構與其他視覺 Transformer 的比較

下表顯示了上述代表性同期工作與CvT在位置編碼的必要性、token嵌入型別、投影型別以及主幹中的Transformer結構方面的關鍵差異。

模型	需要位置編碼 (PE)	Token 嵌入	注意力投影	分層Transformer
ViT[1], DeiT [3]	是	不重疊	線性	否
CPVT[4]	否（帶 PE 生成器）	不重疊	線性	否
TNT[5]	是	不重疊 (Patch + 畫素)	線性	否
T2T[6]	是	重疊（拼接）	線性	部分（Token化）
PVT[7]	是	不重疊	空間縮減	是
CvT[2]	否	重疊（卷積）	卷積	是

主要亮點

CvT 能夠實現卓越效能和計算效率的四個主要亮點如下：

包含新的卷積Token嵌入的Transformer層次結構。
利用卷積投影的卷積Transformer塊。
由於卷積引入的內建區域性上下文結構，移除了位置編碼。
與其他視覺Transformer架構相比，引數更少，浮點運算 (FLOPs) 更低。

PyTorch 實現

動手時間到了！讓我們來探索如何在 PyTorch 中實現 CvT 架構的每個主要模組，如官方實現[8]所示。

匯入所需庫

from collections import OrderedDict
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
from einops.layers.torch import Rearrange

卷積投影的實現

def _build_projection(self, dim_in, dim_out, kernel_size, padding, stride, method):
    if method == "dw_bn":
        proj = nn.Sequential(
            OrderedDict(
                [
                    (
                        "conv",
                        nn.Conv2d(
                            dim_in,
                            dim_in,
                            kernel_size=kernel_size,
                            padding=padding,
                            stride=stride,
                            bias=False,
                            groups=dim_in,
                        ),
                    ),
                    ("bn", nn.BatchNorm2d(dim_in)),
                    ("rearrage", Rearrange("b c h w -> b (h w) c")),
                ]
            )
        )
    elif method == "avg":
        proj = nn.Sequential(
            OrderedDict(
                [
                    (
                        "avg",
                        nn.AvgPool2d(
                            kernel_size=kernel_size,
                            padding=padding,
                            stride=stride,
                            ceil_mode=True,
                        ),
                    ),
                    ("rearrage", Rearrange("b c h w -> b (h w) c")),
                ]
            )
        )
    elif method == "linear":
        proj = None
    else:
        raise ValueError("Unknown method ({})".format(method))

    return proj

該方法接受多個與卷積層相關的引數（如輸入和輸出維度、核大小、填充、步幅和方法），並根據指定的方法返回一個投影塊。

如果方法是dw_bn（帶批次歸一化的深度可分離），它會建立一個由深度可分離卷積層後跟批次歸一化組成的順序塊，並重新排列維度。
如果方法是avg（平均池化），它會建立一個帶有平均池化層後跟維度重新排列的順序塊。
如果方法是linear，它返回None，表示未應用投影。

維度的重新排列透過`Rearrange`操作執行，該操作會重塑輸入張量。然後返回生成的投影塊。

卷積Token嵌入的實現

class ConvEmbed(nn.Module):
    def __init__(
        self, patch_size=7, in_chans=3, embed_dim=64, stride=4, padding=2, norm_layer=None
    ):
        super().__init__()
        patch_size = to_2tuple(patch_size)
        self.patch_size = patch_size

        self.proj = nn.Conv2d(
            in_chans, embed_dim, kernel_size=patch_size, stride=stride, padding=padding
        )
        self.norm = norm_layer(embed_dim) if norm_layer else None

    def forward(self, x):
        x = self.proj(x)

        B, C, H, W = x.shape
        x = rearrange(x, "b c h w -> b (h w) c")
        if self.norm:
            x = self.norm(x)
        x = rearrange(x, "b (h w) c -> b c h w", h=H, w=W)

        return x

此程式碼定義了一個ConvEmbed模組，用於對輸入影像執行逐塊嵌入。

__init__方法使用各種引數初始化模組，例如patch_size（影像塊的大小）、in_chans（輸入通道數）、embed_dim（嵌入塊的維度）、stride（卷積操作的步幅）、padding（卷積操作的填充）和norm_layer（可選的歸一化層）。
在建構函式中，使用指定的引數（包括影像塊大小、輸入通道、嵌入維度、步幅和填充）建立2D卷積層（nn.Conv2d）。此卷積層分配給self.proj。
如果提供了歸一化層，則會建立具有embed_dim通道的歸一化層例項，並將其分配給self.norm。
前向方法接收輸入張量x，並使用self.proj應用卷積操作。使用rearrange函式重塑輸出以平鋪空間維度。如果存在歸一化層，則將其應用於平鋪表示。最後，將張量重塑回原始空間維度並返回。

總而言之，該模組旨在對影像進行逐塊嵌入，其中每個塊透過卷積層獨立處理，並對嵌入特徵應用可選的歸一化。

視覺Transformer塊的實現

class VisionTransformer(nn.Module):
    """Vision Transformer with support for patch or hybrid CNN input stage"""

    def __init__(
        self,
        patch_size=16,
        patch_stride=16,
        patch_padding=0,
        in_chans=3,
        embed_dim=768,
        depth=12,
        num_heads=12,
        mlp_ratio=4.0,
        qkv_bias=False,
        drop_rate=0.0,
        attn_drop_rate=0.0,
        drop_path_rate=0.0,
        act_layer=nn.GELU,
        norm_layer=nn.LayerNorm,
        init="trunc_norm",
        **kwargs,
    ):
        super().__init__()
        self.num_features = self.embed_dim = embed_dim

        self.rearrage = None

        self.patch_embed = ConvEmbed(
            patch_size=patch_size,
            in_chans=in_chans,
            stride=patch_stride,
            padding=patch_padding,
            embed_dim=embed_dim,
            norm_layer=norm_layer,
        )

        with_cls_token = kwargs["with_cls_token"]
        if with_cls_token:
            self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        else:
            self.cls_token = None

        self.pos_drop = nn.Dropout(p=drop_rate)
        dpr = [
            x.item() for x in torch.linspace(0, drop_path_rate, depth)
        ]  # stochastic depth decay rule

        blocks = []
        for j in range(depth):
            blocks.append(
                Block(
                    dim_in=embed_dim,
                    dim_out=embed_dim,
                    num_heads=num_heads,
                    mlp_ratio=mlp_ratio,
                    qkv_bias=qkv_bias,
                    drop=drop_rate,
                    attn_drop=attn_drop_rate,
                    drop_path=dpr[j],
                    act_layer=act_layer,
                    norm_layer=norm_layer,
                    **kwargs,
                )
            )
        self.blocks = nn.ModuleList(blocks)

        if self.cls_token is not None:
            trunc_normal_(self.cls_token, std=0.02)

        if init == "xavier":
            self.apply(self._init_weights_xavier)
        else:
            self.apply(self._init_weights_trunc_normal)

        def forward(self, x):
            x = self.patch_embed(x)
            B, C, H, W = x.size()

            x = rearrange(x, "b c h w -> b (h w) c")

            cls_tokens = None
            if self.cls_token is not None:
                cls_tokens = self.cls_token.expand(B, -1, -1)
                x = torch.cat((cls_tokens, x), dim=1)

            x = self.pos_drop(x)

            for i, blk in enumerate(self.blocks):
                x = blk(x, H, W)

            if self.cls_token is not None:
                cls_tokens, x = torch.split(x, [1, H * W], 1)
            x = rearrange(x, "b (h w) c -> b c h w", h=H, w=W)

            return x, cls_tokens

此程式碼定義了一個視覺Transformer模組。以下是程式碼的簡要概述：

初始化： `VisionTransformer`類使用定義模型架構的各種引數進行初始化，例如補丁大小、嵌入維度、層數、注意力頭數、 dropout 率等。
補丁嵌入： 該模型包含一個補丁嵌入層（`patch_embed`），它透過將輸入影像劃分為不重疊的補丁並使用卷積進行嵌入來處理輸入影像。
Transformer 塊： 該模型由一堆 Transformer 塊（`Block`）組成。塊的數量由深度引數決定。每個塊都包含多頭自注意力機制和前饋神經網路。
分類Token： 可選地，模型可以包含一個可學習的分類Token（`cls_token`），附加到輸入序列中。該Token用於分類任務。
隨機深度： 隨機深度應用於Transformer塊，其中在訓練期間跳過隨機子集塊以改善正則化。這由`drop_path_rate`引數控制。
權重初始化： 模型權重使用截斷正態分佈（`trunc_norm`）或 Xavier 初始化（`xavier`）進行初始化。
前向方法： 前向方法透過補丁嵌入處理輸入，重新排列維度，如果存在則新增分類Token，應用dropout，然後透過Transformer塊堆疊傳遞資料。最後，將輸出重新排列回原始形狀，並將分類Token（如果存在）與序列的其餘部分分離，然後返回輸出。

卷積視覺Transformer塊（Transformer 層次結構）的實現

class ConvolutionalVisionTransformer(nn.Module):
    def __init__(
        self,
        in_chans=3,
        num_classes=1000,
        act_layer=nn.GELU,
        norm_layer=nn.LayerNorm,
        init="trunc_norm",
        spec=None,
    ):
        super().__init__()
        self.num_classes = num_classes

        self.num_stages = spec["NUM_STAGES"]
        for i in range(self.num_stages):
            kwargs = {
                "patch_size": spec["PATCH_SIZE"][i],
                "patch_stride": spec["PATCH_STRIDE"][i],
                "patch_padding": spec["PATCH_PADDING"][i],
                "embed_dim": spec["DIM_EMBED"][i],
                "depth": spec["DEPTH"][i],
                "num_heads": spec["NUM_HEADS"][i],
                "mlp_ratio": spec["MLP_RATIO"][i],
                "qkv_bias": spec["QKV_BIAS"][i],
                "drop_rate": spec["DROP_RATE"][i],
                "attn_drop_rate": spec["ATTN_DROP_RATE"][i],
                "drop_path_rate": spec["DROP_PATH_RATE"][i],
                "with_cls_token": spec["CLS_TOKEN"][i],
                "method": spec["QKV_PROJ_METHOD"][i],
                "kernel_size": spec["KERNEL_QKV"][i],
                "padding_q": spec["PADDING_Q"][i],
                "padding_kv": spec["PADDING_KV"][i],
                "stride_kv": spec["STRIDE_KV"][i],
                "stride_q": spec["STRIDE_Q"][i],
            }

            stage = VisionTransformer(
                in_chans=in_chans,
                init=init,
                act_layer=act_layer,
                norm_layer=norm_layer,
                **kwargs,
            )
            setattr(self, f"stage{i}", stage)

            in_chans = spec["DIM_EMBED"][i]

        dim_embed = spec["DIM_EMBED"][-1]
        self.norm = norm_layer(dim_embed)
        self.cls_token = spec["CLS_TOKEN"][-1]

        # Classifier head
        self.head = (
            nn.Linear(dim_embed, num_classes) if num_classes > 0 else nn.Identity()
        )
        trunc_normal_(self.head.weight, std=0.02)

    def forward_features(self, x):
        for i in range(self.num_stages):
            x, cls_tokens = getattr(self, f"stage{i}")(x)

        if self.cls_token:
            x = self.norm(cls_tokens)
            x = torch.squeeze(x)
        else:
            x = rearrange(x, "b c h w -> b (h w) c")
            x = self.norm(x)
            x = torch.mean(x, dim=1)

        return x

    def forward(self, x):
        x = self.forward_features(x)
        x = self.head(x)

        return x

此程式碼定義了一個名為ConvolutionalVisionTransformer的 PyTorch 模組。

該模型由多個階段組成，每個階段由一個VisionTransformer類的例項表示。
每個階段都有不同的配置，例如spec字典中指定的補丁大小、步幅、深度、頭部數量等。
forward_features方法處理透過所有階段的輸入x，並聚合最終表示。
該類具有一個分類器頭，用於執行線性變換以產生最終輸出。
forward方法呼叫forward_features，然後將結果傳遞給分類器頭。
視覺Transformer階段按順序命名為 stage0、stage1 等，每個階段都是VisionTransformer類的一個例項，形成一個Transformer層次結構。

恭喜！現在您已經知道如何在PyTorch中實現CvT架構了。您可以在此處檢視CvT架構的完整程式碼。

試用

如果您想使用 CvT，而不想深入瞭解其 PyTorch 實現的複雜細節，您可以輕鬆地利用 Hugging Face `transformers` 庫。操作如下：

pip install transformers

您可以在此處找到 CvT 模型的文件。

使用方法

以下是使用 CvT 模型將 COCO 2017 資料集中的影像分類為 1000 個 ImageNet 類別之一的方法：

from transformers import AutoFeatureExtractor, CvtForImageClassification
from PIL import Image
import requests

url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)

feature_extractor = AutoFeatureExtractor.from_pretrained("microsoft/cvt-13")
model = CvtForImageClassification.from_pretrained("microsoft/cvt-13")

inputs = feature_extractor(images=image, return_tensors="pt")
outputs = model(**inputs)
logits = outputs.logits
# model predicts one of the 1000 ImageNet classes
predicted_class_idx = logits.argmax(-1).item()
print("Predicted class:", model.config.id2label[predicted_class_idx])

參考文獻

< > 在 GitHub 上更新

社群計算機視覺課程