卷积视觉Transformer (CvT)

在本节中，我们将深入探讨卷积视觉Transformer (CvT)，它是视觉Transformer (ViT) 的变体[1]，广泛用于计算机视觉中的图像分类任务。

回顾

在深入探讨 CvT 之前，让我们简要回顾一下之前章节中介绍的 ViT 架构，以便更好地理解 CvT 架构。ViT 将每个图像分解成一系列固定长度的标记（即不重叠的补丁），然后应用多个标准 Transformer 层，包括多头自注意力和位置感知前馈模块 (FFN)，以对全局关系进行建模，以用于分类。

概述

卷积视觉Transformer (CvT) 模型是在 CvT：将卷积引入视觉Transformer 中提出的[2]，由吴海平、肖斌、诺埃尔·科德拉、刘梦晨、戴夕洋、袁路和张磊提出。CvT 采用了 CNN 的所有优势：局部感受野、共享权重和空间下采样，以及移位、缩放、失真不变性，同时保留了 Transformer 的优点：动态注意力、全局上下文融合和更好的泛化能力。与 ViT 相比，CvT 在保持计算效率的同时，取得了优越的性能。此外，由于卷积引入了内置的局部上下文结构，CvT 不再需要位置嵌入，这使其有可能适应各种需要可变输入分辨率的视觉任务。

架构

CvT 架构 (a) 整体架构，展示了卷积标记嵌入层所促成的分层多级结构。(b) 卷积Transformer 块的详细信息，其中包含卷积投影作为第一层。 [2]

上面 CvT 架构的图像说明了 3 级管道的主要步骤。CvT 的核心是将两种基于卷积的操作融合到视觉Transformer 架构中

卷积标记嵌入：想象将输入图像分成重叠的补丁，将它们重塑成标记，然后将它们馈送到卷积层。这减少了标记的数量（就像下采样图像中的像素），同时提高了它们的特征丰富度，类似于传统的 CNN。与其他 Transformer 不同的是，我们跳过将预定义的位置信息添加到标记中，完全依赖于卷积运算来捕获空间关系。

投影层 (a) ViT 中的线性投影。(b) 卷积投影。(c) 压缩卷积投影（CvT 中的默认设置）。 [2]

卷积Transformer 块：CvT 中的每个阶段都包含一系列这样的块。在这里，我们使用深度可分离卷积（卷积投影）来处理自注意力模块的“查询”、“键”和“值”组件，而不是 ViT 中的常规线性投影，如上图所示。这保持了 Transformer 的优点，同时提高了效率。请注意，“分类标记”（用于最终预测）仅在最后阶段添加。最后，一个标准的全连接层分析最终的分类标记以预测图像类别。

CvT 架构与其他视觉Transformer 的比较

下表显示了上述代表性并发工作与 CvT 之间在骨干网络中位置编码的必要性、标记嵌入的类型、投影的类型和 Transformer 结构方面的关键区别。

模型	需要位置编码 (PE)	标记嵌入	用于注意力的投影	分层Transformer
ViT[1]、DeiT [3]	是	不重叠	线性	否
CPVT[4]	否（使用 PE 生成器）	不重叠	线性	否
TNT[5]	是	非重叠（Patch + Pixel）	线性	否
T2T[6]	是	重叠（拼接）	线性	部分（分词）
PVT[7]	是	不重叠	空间降维	是
CvT[2]	否	重叠（卷积）	卷积	是

主要亮点

CvT 为实现更优越的性能和计算效率而采用的四个主要亮点如下

Transformer 层次结构包含一个新的卷积 Token 嵌入。
利用卷积投影的卷积 Transformer 模块。
由于卷积引入的内置局部上下文结构，移除位置编码。
与其他视觉 Transformer 架构相比，参数更少，FLOPs（每秒浮点运算次数）更低。

PyTorch 实现

现在开始动手操作！让我们探索如何使用 PyTorch 对官方实现[8]中所示的 CvT 架构的每个主要模块进行编码。

导入所需库

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（带批归一化的深度可分离卷积），则它会创建一个由深度可分离卷积层后跟批归一化组成的 Sequential 块，并重新排列维度。
如果方法是avg（平均池化），则它会创建一个包含平均池化层和维度重排的 Sequential 块。
如果方法是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（归一化层，可选）之类的参数初始化模块。
在构造函数中，会使用指定参数创建一个二维卷积层（nn.Conv2d），包括块大小、输入通道、嵌入维度、步长和填充。此卷积层分配给self.proj。
如果提供了归一化层，则会使用embed_dim通道创建一个归一化层实例，并将它分配给self.norm。
forward 方法接受一个输入张量 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）进行初始化。
forward 方法：forward 方法通过块嵌入处理输入，重新排列维度，添加分类 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 架构的完整代码here.

试一试

如果您不想深入了解 CvT 的复杂 PyTorch 实现，就可以通过利用 Hugging Face 的transformers库轻松实现。方法如下

pip install transformers

您可以在此处找到 CvT 模型的文档here.

用法

以下是使用 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 上更新

社区计算机视觉课程

卷积视觉Transformer (CvT)

回顾

概述

架构

CvT 架构与其他视觉Transformer 的比较

主要亮点

PyTorch 实现

试一试

用法

参考文献