社区计算机视觉课程文档

卷积视觉 Transformer (CvT)

Hugging Face's logo
加入 Hugging Face 社区

并获取增强的文档体验

开始使用

卷积视觉 Transformer (CvT)

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

概括

在深入研究 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 中提出。CvT 结合了 CNN 的所有优点:局部感受野、权重共享和空间子采样,以及平移、缩放、失真不变性,同时保留了 Transformer 的优点:动态注意力、全局上下文融合和更好的泛化能力。与 ViT 相比,CvT 在保持计算效率的同时实现了卓越的性能。此外,由于卷积引入的内置局部上下文结构,CvT 不再需要位置嵌入,这使其在适应需要可变输入分辨率的各种视觉任务方面具有潜在优势。

架构

CvT Architecture (a) 总体架构,展示了卷积 Token 嵌入层促进的分层多阶段结构。(b) 卷积 Transformer 块的详细信息,其中包含作为第一层的卷积投影。 [2]

上面的 CvT 架构图示说明了 3 阶段流程的主要步骤。CvT 的核心是将两个基于卷积的操作融入到视觉 Transformer 架构中

  • 卷积 Token 嵌入:想象一下将输入图像分割成重叠的图像块,将它们重塑为 token,然后将它们馈送到卷积层。这减少了 token 的数量(类似于下采样图像中的像素),同时提高了它们的特征丰富性,类似于传统的 CNN。与其他 Transformer 不同,我们跳过向 token 添加预定义的位置信息,而仅依靠卷积运算来捕获空间关系。

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

  • 卷积 Transformer 块:CvT 中的每个阶段都包含一堆这样的块。在这里,我们没有使用 ViT 中常用的线性投影,而是使用深度可分离卷积(卷积投影)来处理自注意力模块的“query”、“key”和“value”组件,如上图所示。这既保留了 Transformer 的优势,又提高了效率。请注意,“分类 token”(用于最终预测)仅在最后阶段添加。最后,一个标准的完全连接层分析最终的分类 token 以预测图像类别。

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

下表显示了上述代表性并行工作和 CvT 之间在位置编码的必要性、token 嵌入类型、投影类型以及骨干网络中的 Transformer 结构方面的主要差异。

模型 需要位置编码 (PE) Token 嵌入 注意力投影 分层 Transformer
ViT[1], DeiT [3] 非重叠 线性
CPVT[4] 否(带 PE 生成器) 非重叠 线性
TNT[5] 非重叠(图像块 + 像素) 线性
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(深度可分离卷积,带批归一化),它将创建一个 Sequential 块,该块由深度可分离卷积层组成,后跟批归一化并重新排列维度。

  • 如果方法是 avg(平均池化),它将创建一个 Sequential 块,其中包含一个平均池化层,然后重新排列维度。

  • 如果方法是 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

  • forward 方法接受输入张量 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) 组成。块的数量由 depth 参数确定。每个块都包含多头自注意力机制和前馈神经网络。

  • 分类 Token:可选地,该模型可以包含一个可学习的分类 token (cls_token),该 token 附加到输入序列。此 token 用于分类任务。

  • 随机深度:随机深度应用于 Transformer 块,其中在训练期间跳过随机子集的块以改善正则化。这由 drop_path_rate 参数控制。

  • 权重初始化:模型权重使用截断正态分布 (trunc_norm) 或 Xavier 初始化 (xavier) 进行初始化。

  • Forward 方法:forward 方法通过图像块嵌入处理输入,重新排列维度,如果存在则添加分类 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 数据集的图像分类为 1,000 个 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 上更新