Vision transformer

太长不看版

Vision Transformer将CV和NLP领域知识结合起来,对原始图片进行分块,展平成序列,输入进原始Transformer模型的编码器Encoder部分,最后接入一个全连接层对图片进行分类。在大型数据集上表现超过了当前SOTA模型

前言

当前Transformer模型被大量应用在NLP自然语言处理当中,而在计算机视觉领域,Transformer的注意力机制attention也被广泛应用,比如Se模块,CBAM模块等等注意力模块,这些注意力模块能够帮助提升网络性能。而我们的工作展示了不需要依赖CNN的结构,也可以在图像分类任务上达到很好的效果,并且也十分适合用于迁移学习。

这里的代码引用自 https://github.com/lucidrains/vit-pytorch,大家有兴趣也可以跑跑demo。

方法

首先结构上,我们采取的是原始Transformer模型,方便开箱即用。

如果对Transformer模型不太了解的可以参考这篇文章 解析Transformer模型

整体结构如下
Vision Transformer整体架构

数据处理部分

原始输入的图片数据是 H x W x C,我们先对图片作分块,再进行展平。假设每个块的长宽为(P, P),那么分块的数目为
N = H ∗ W / ( P ∗ P ) N = H * W /(P*P) N=HW/(PP)
然后对每个图片块展平成一维向量,每个向量大小为
P ∗ P ∗ C P*P*C PPC
总的输入变换为
N × ( P 2 ⋅ C ) N \times (P^2·C) N×(P2C)
这里的代码如下:

x = rearrange(img, 'b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=p, p2=p)

它使用的是一个einops的拓展包,完成了上述的变换工作

Patch Embedding

接着对每个向量都做一个线性变换(即全连接层),压缩维度为D,这里我们称其为 Patch Embedding。

在代码里是初始化一个全连接层,输出维度为dim,然后将分块后的数据输入

self.patch_to_embedding = nn.Linear(patch_dim, dim)

# forward前向代码
x = rearrange(img, 'b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=p, p2=p)
x = self.patch_to_embedding(x)

Positional Encoding

还记得在解析Transformer那篇文章内有说过,原始的Transformer引入了一个 Positional encoding 来加入序列的位置信息,同样在这里也引入了pos_embedding,是用一个可训练的变量替代。

self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))

文章也提供了可视化图
Pos Encoding可视化

很有意思的是这里第二个维度多加了个1。下面会有讲到

class_token

这里我们再来仔细看上图的一个结构
额外插入一个embedding
假设我们按照论文切成了9块,但是在输入的时候变成了10个向量。这是人为增加的一个向量。

因为传统的Transformer采取的是类似seq2seq编解码的结构
而ViT只用到了Encoder编码器结构,缺少了解码的过程,假设你9个向量经过编码器之后,你该选择哪一个向量进入到最后的分类头呢?因此这里作者给了额外的一个用于分类的向量,与输入进行拼接。同样这是一个可学习的变量。

具体操作如下

# 假设dim=128,这里shape为(1, 1, 128)
self.cls_token = nn.Parameter(torch.randn(1, 1, dim))

# forward前向代码
# 假设batchsize=10,这里shape为(10, 1, 128)
cls_tokens = repeat(self.cls_token, '() n d -> b n d', b=b)
# 跟前面的分块为x(10,64, 128)的进行concat
# 得到(10, 65, 128)向量
x = torch.cat((cls_tokens, x), dim=1)

知道这个操作,我们也就能明白为什么前面的pos_embedding的第一维也要加1了,后续将pos_embedding也加入到x

 x += self.pos_embedding[:, :(n + 1)]

分类

分类头很简单,加入了LayerNorm和两层全连接层实现的,采用的是GELU激活函数。
代码如下

self.mlp_head = nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, mlp_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(mlp_dim, num_classes)
        )

最终分类我们只取第一个,也就是用于分类的token,输入到分类头里,得到最后的分类结果

self.to_cls_token = nn.Identity()
# forward前向部分
x = self.transformer(x, mask)
x = self.to_cls_token(x[:, 0])
return self.mlp_head(x)

可以看到整个流程是非常简单的,下面是ViT的整体代码

class ViT(nn.Module):
    def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, channels=3, dropout=0.,
                 emb_dropout=0.):
        super().__init__()
        assert image_size % patch_size == 0, 'Image dimensions must be divisible by the patch size.'
        num_patches = (image_size // patch_size) ** 2
        patch_dim = channels * patch_size ** 2
        assert num_patches > MIN_NUM_PATCHES, f'your number of patches ({num_patches}) is way too small for attention to be effective (at least 16). Try decreasing your patch size'

        self.patch_size = patch_size

        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
        self.patch_to_embedding = nn.Linear(patch_dim, dim)
        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
        self.dropout = nn.Dropout(emb_dropout)

        self.transformer = Transformer(dim, depth, heads, mlp_dim, dropout)

        self.to_cls_token = nn.Identity()

        self.mlp_head = nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, mlp_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(mlp_dim, num_classes)
        )

    def forward(self, img, mask=None):
        p = self.patch_size

        x = self.patch_to_embedding(x)
        b, n, _ = x.shape

        cls_tokens = repeat(self.cls_token, '() n d -> b n d', b=b)
        x = torch.cat((cls_tokens, x), dim=1)
        x += self.pos_embedding[:, :(n + 1)]
        x = self.dropout(x)

        x = self.transformer(x, mask)

        x = self.to_cls_token(x[:, 0])
        return self.mlp_head(x)

实验部分

与Transformer一样,ViT也有规模不一样的模型设置,如下图所示
不同规模ViT的配置
可以看到整体模型还是挺大的,而经过大数据集的预训练后,性能也超过了当前CNN的一些SOTA结果
预训练后的ViT性能比较
另外作者还给了注意力观察得到的图片块,我的一点猜想是可能有利于对神经网络可解释性的研究。
注意力可视化

总结

继DETR后,这又是一个CV和NLP结合的工作。思想非常的朴素简单,就是拿最原始的Transformer模型来做图像分类。现有的性能还需要大量的数据来训练,期待后续工作对ViT做一些改进,降低其训练时间和所需数据量,让人人都能玩得起ViT!

  • 57
    点赞
  • 241
    收藏
    觉得还不错? 一键收藏
  • 18
    评论
Vision Transformer(ViT)是一种基于Transformer架构的深度学习模型,用于处理计算机视觉任务。它将图像分割成一系列的图像块,并将每个图像块作为输入序列传递给Transformer编码器。每个图像块通过一个线性投影层转换为向量表示,并与位置嵌入向量相结合,然后输入到Transformer编码器中进行处理。Transformer编码器由多个自注意力层和前馈神经网络层组成,用于学习图像中的全局和局部特征。最后,通过一个线性分类器对编码器的输出进行分类。 Vision Transformer的优点是能够在没有使用传统卷积神经网络的情况下,实现对图像的高质量特征提取和分类。它在一些计算机视觉任务上取得了与传统方法相媲美甚至更好的性能,例如图像分类、目标检测和语义分割等任务。 以下是一个使用Vision Transformer进行图像分类的示例代码[^1]: ```python import torch import torch.nn as nn from torchvision import transforms from torchvision.models import resnet50 from vit_pytorch import ViT # 加载预训练的Vision Transformer模型 model = ViT( image_size = 224, patch_size = 16, num_classes = 1000, dim = 768, depth = 12, heads = 12, mlp_dim = 3072, dropout = 0.1, emb_dropout = 0.1 ) # 加载预训练的权重 model.load_from('vit_weights.pth') # 图像预处理 transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) # 加载图像并进行预处理 image = Image.open('image.jpg') image = transform(image).unsqueeze(0) # 使用Vision Transformer进行图像分类 output = model(image) _, predicted_class = torch.max(output, 1) # 输出预测结果 print('Predicted class:', predicted_class.item()) ```
评论 18
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值