社区计算机视觉课程文档

卷积视觉Transformer (CvT)

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

卷积视觉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]所示。

  1. 导入所需库
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
  1. 卷积投影的实现
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`操作执行,该操作会重塑输入张量。然后返回生成的投影块。

  1. 卷积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函数重塑输出以平铺空间维度。如果存在归一化层,则将其应用于平铺表示。最后,将张量重塑回原始空间维度并返回。

总而言之,该模块旨在对图像进行逐块嵌入,其中每个块通过卷积层独立处理,并对嵌入特征应用可选的归一化。

  1. 视觉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(如果存在)与序列的其余部分分离,然后返回输出。

  1. 卷积视觉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 上更新