下面借这篇blog记录一些阅读笔记，如果问题，恳请指出。
文章目录
前言
paper原文：AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE
一开始，Transformer在Attention is all you need这篇paper中提出，解决的自然语言中的序列问题，也就是将自然语言的word变成一个sequence问题，但是有效的解决了RNN的无法并行处理的问题，并且其可以考虑全局的咨询，而self-attention主要工作就是矩阵运算，所以可以使用硬件加速。
而这篇paper是将Transformer应用在图像识别任务上，用paper题目的意思就是，一张图片可以切分为多个16x16的图片，而以分辨率为224的图片来说，可以切分14份这个的16x16的图片，每一张图片可以称其为一个patch。下面会将Vision Transformer称为vit模型。
为了与原来的Transformer足够的类似，vit模型模型的流程大概如下，其会将输入的图片切分为多个patch，将这些patch加以处理之后会通过Transformer的Encoder进行处理，然后再通过一个分类器的head，得到每张图片的预测结果，详细结构见下文。
1.Vision Transformer结构
基于上面的简要认识，现在具体的谈谈其中的细节。
Vision Transformer的整体结构如图所示：
1.1 Patch and Positional Encoding
假设我们需要处理的图像的分辨率是224x224，可以有图看见，我们会将原始图像切分为一份份的patch，这些patch的分辨率大小有16x16，或者是32x32，现在假设切分大小为16x16的大小，那么原始图像的224x224大小可以切分出14x14个这样的pacth，也就是196个patch。
那么这是如何进行切分操作的呢？其实是通过一个普通的卷积操作con2v实现，通过一个大小为patch_size的卷积核，stride同样为patch_size，进行卷积处理，就可以实现以上的流程。而我们知道，卷积操作之后是得到了一个数值，所以我们得到了14x14这样的特征矩阵，也就是196个patch。此外在代码中，模型使用了768个这样的卷积核。
进行了以上卷积操作后，得到的196个patch其实就是Transformer的sequence输入，满足输入的条件的。768个channels的输出就相当于经过卷积特征简单提取后每个patch所携带的图像信息。
但是，这样还不够，将这196个patch从1开始编号，然后在第0位置增加一个patch，将其称为class token，也就是假设其可以携带这个图片的信息，在最后的输出时会将其提取出来，判断这张图片属于哪一类。而由于之前每一个patch所携带的信息的dim都是768，所以增加的这个class token的dim也是768的。所以，现在的patch大小一共为197.
但是，在上次我总结Transformer笔记的时候，传送门：学习笔记——Transformer结构的完整介绍。有提到，直接在做self-attention之前，我们需要增加一些位置信息，所以进行位置编码操作，也就是Positional Encoding。在vit模型中，也不例外。在参考代码的操作中，vit模型是使用了nn.Parameter来定义pos_embed，也就是表示pos_embed的值是可以被训练的，其是作为module的可训练参数，即加入到parameter()这个迭代器中去。然后将位置编码信息pos_embed与上诉所说的patch信息相加，将相当于输入已经融合了位置信息了。当然，pos_embed的shape需要一致才能做residual add操作。
1.2 Transformer Encoder Block
现在的patch添加了Positional Encoding信息之后，就变成了如图所示的Embedded Patches。接下来就是将Embedded Patches输入到Transformer 中了。
其实，在vit模型中的Transformer Encoder就是原本Transformer Encoder，结构上基本是一样的，所以paper原文也说了，他们对原始的Transformer作出了最大的保留，尽量不改变模型结构。换一句话来说，vit模型就是使用了Transformer的Encoder结构实现了图像的分类。
这里再做一个简单的介绍，在vit模型中，Embedded Patches首先进行Layer Norm处理，然后进行多头的Self-allention操作，再作一个residual add操作。接下来，再进行Layer Norm处理，然后通过一个多层感知机也就是MLP层（代码实现中其实就是两层全连接层，激活函数使用的是GELU激活函数。其中第一层全连接层将768维度的信息扩充了4倍，也就是变成3072dim，然后第二层全连接层在将其缩回原来的768维），然后再做一个residual add操作。这个就是Transformer Encoder Block的整个流程，然后将其进行不断的堆叠，在代码中对其堆叠了12次。
这里还需要注意，对于Encoder堆叠的block结构，使用的是DropPath操作，而在MLP层中的，使用的是Dropout操作。
具体的Block结构在pytorch中如图所示：
1.3 Classifier head
在Transformer Encoder中，不断的对Block进行堆叠，其实也可以看成是一个特征提取的流程，同样的，在特征提取完之后，就通过分类器，进行最后的输出结果的预测。
但是需要注意，这里进行分类预测的不是所以的patch，而是197个patch中的第一个patch，也就是那个第0个位置的class token，现在将其提出出来进行最后的分类预测。由上诉我们见到，一个patch所包含的信息维度是768，所以现在只需要简单的将这768维的信息变成一个dim = class_nums就可以了。
以上便是vit模型操作过程中的整个流程。
2.Vision Transformer细节
第一次自己归纳一篇paper的内容，还是感觉有很多的细节的，只能以这么一节来记录一下。
2.1 performance
paper中一共提出了3中vit模型，分别是基本版本base，大版本large，还有一个超大版本huge，分别对block堆叠次数layers，patch所携带的信息dims，多层感知机的hidden size（也就是扩充倍数），还有Multied-head self-attention中heads的个数进行改动，如图所示：
而在这三个版本之外，还对patch_size进行了改动，分为16x16与32x32两个版本。由于16x16分得更小，所以sequence的长度会更长，所以参数量相比32x32的patch_size会有所增长。
作者将这几种模型之间进行比较，如图所示：
这张图显示了当在ImageNet、ImageNet-21k或JFT300M上进行预训练时，Vision Transformer在各种数据集上的精度大小。其中，在一些比较大的数据集上进行预训练对于vit模型来说是及其重要的，如果没有预训练，vit模型的效果是不怎么好的。而可以看见，我标注了一句话，这些模型都是在384的分辨率进行微调的，而且paper中也有提到，在高分辨率上进行微调是非常有帮助的。
It is often beneficial to fine-tune at higher resolution than pre-training
对于上图，可以发现，随着预训练数据集的不断扩大，模型的效果也越来越好。
而如果想自己从头开始训练模型时，一个比较强的正则化是及其重要的，但是这需要太多的时间，使用一般进行迁移学习然后进行fine-tune就好。
此外，paper还对深度与宽度进行调整对比，如图所示：
可以看到，扩展深度带来了最大的改进，直到64层都清晰可见。然而，对于扩展宽度的收益递减在16层之后已经可见。也就是说增加深度要比增加宽度要好。增加宽度带来的收益不会太大。所以，如果考虑缩放，那么缩放应该强调深度而不是宽度。
此外，减小patch_size导致的sequence变长会增加模型的鲁棒性，所以这就是为什么上图显示，patch_size为16x16的模型为什么要比32x32的效果更好，虽然这增加了计算量，但是提高的效果。
2.2 hybrid architecture
其中paper中还提到了一个Hybrid Architecture，也就是cnn与transformer的混合结构。
与vit模型结构不同，输入到transformer encoder的输入，可以不是images patch，还可以是经过cnn特征提取到的sequence，也可以称其为特征矩阵，这可以看成是一个特殊的patch，不过这些patch所携带的信息dim=1，已经被高度的抽象。也就是将cnn提取到的特征矩阵看成是一连串的patch，所以每个patch只包含一个数值，也就是携带一个信息。这意味着通过简单地展平特征图的空间维度并投影到Transformer维度来获得输入序列。
当然，提取到的特征矩阵同样需要进行position embedding处理，需要添加位置信息。
paper分别将Hybrid Architecture，与vit模型结构，再与ResNet模型进行对比：
可以看见，在一个庞大的数据集预训练之后的vit模型要比resnet效果要好。而对于混合模型来说，在训练比较少的epcoh的时候其是有优势的，但当训练的epoch足够的时候，混合模型是比不过正常的vit模型的。
paper总结如下：
1）transformer少使用了2-4倍的计算去获得了与Resnet的相同性能
2）在混合模型较小的计算预算是优于vit模型，但是对于训练次数够多时，这种差异减少，也就是说，利用cnn来局部特行处理的辅助任何大小transformer的效果不一定比只使用transformer要好。
3）transformer在没有在尝试的范围内饱和，其拓展性在未来可能更好。
2.3 computational cost
此外，在paper中还对比的预测时间与消耗量。主要与ResNet来对比，如图所示：
可以看见，无论是每秒可以分类的图像数量还是内容的消耗，vit模型都要比ResNet要好。
但是，经过如此，在训练我们自己的数据集时，我觉得vit模型训练还是需要比较大的内存的。
paper中还提到了两个点：
1）对于最大分辨率下的最大模型，理论上ViT随图像尺寸的双二次缩放几乎没有开始发生。
2）每个模型可以放入核心的最大批量的数量越大就越适合扩展到大型数据集。
不过我不太懂这两句话的含义，希望有大佬可以指点一下。
2.4 positional embedding
在上诉我们提到的位置编码，其实有多种的实现方式。
- Providing no positional information：不处理
- 1-dimensional positional embedding：按顺序，一维坐标表示
- 2-dimensional positional embedding：按顺序，二维坐标表示
- Relative positional embeddings：考虑每个patch的相对距离
结果如图所示：
没有位置嵌入的模型和有位置嵌入的模型之间的性能有很大的差距，但是编码位置信息的不同方式之间几乎没有差别。
paper中给出了一个猜想：由于Transformer Encoder是在patch级别的输入上运行的，而不是pixel像素级输入，因此如何编码空间信息的差异就不那么重要了。更准确地说，在patch级别输入中，空间维度比原始像素级输入小得多。例如14 × 14，而不是224 × 224，并且对于这些不同的位置编码策略，学习以这种分辨率表示空间关系同样容易。
基于上述讨论，在代码实现中只是简单的采用了第二种1-dimensional positional embedding的方法。
2.5 attention distance
这一点感觉非常神奇且重要。
为了了解 ViT 如何使用self-attention来整合图像中的信息，作者还分析了不同层的注意力权重跨越的平均距离（“attention distance”）。
attention distance类似于cnn的感受野大小。平均注意力距离在较低层的头部之间变化很大，一些头部关注图像的大部分，而其他头部关注查询位置处或附近的小区域。随着深度的增加，所有头部的注意力距离都会增加。在网络的后半部分，大多数heads都与class token有紧密联系。
这对于经典的cnn网络在某一程度上有些类似，在vit模型中，由于使用了多个head，所以有些head会关注局部信息，而有些head会关注全局的信息。而对于cnn网络来说，在比较底层的特征时，其感受野较大，抽象为一个全局特征，类似的在vit模型中此时多数head的attention distance变化巨大；而对于教高层特征是，其感受野较小，只观察局部特征，所以类似的，vit模型此时多数head的attention distance较小。
3.Vision Transformer的pytorch实现
参考代码：
"""
original code from rwightman:
https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
"""
from functools import partial
from collections import OrderedDict
import torch
import torch.nn as nn
# Stochastic Depth随机深度处理方法
def drop_path(x, drop_prob: float = 0., training: bool = False):
"""
Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
'survival rate' as the argument.
"""
if drop_prob == 0. or not training:
return x
keep_prob = 1 - drop_prob
shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
random_tensor.floor_() # binarize
# output: torch.Size([8, 197, 768])
output = x.div(keep_prob) * random_tensor
return output
class DropPath(nn.Module):
"""
Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
"""
def __init__(self, drop_prob=None):
super(DropPath, self).__init__()
self.drop_prob = drop_prob
def forward(self, x):
return drop_path(x, self.drop_prob, self.training)
# 位置编码操作
class PatchEmbed(nn.Module):
"""
2D Image to Patch Embedding
"""
def __init__(self, img_size=224, patch_size=16, in_c=3, embed_dim=768, norm_layer=None):
super().__init__()
img_size = (img_size, img_size)
patch_size = (patch_size, patch_size)
self.img_size = img_size # 224
self.patch_size = patch_size # 16
self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) # (14, 14)
self.num_patches = self.grid_size[0] * self.grid_size[1]
# channels: 3 -> 768
self.proj = nn.Conv2d(in_c, embed_dim, kernel_size=patch_size, stride=patch_size)
self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() # nn.Identity 不处理
def forward(self, x):
# torch.Size([8, 3, 224, 224])
B, C, H, W = x.shape
assert H == self.img_size[0] and W == self.img_size[1], \
f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
# proj: [8, 3, 224, 224] -> [8, 768, 14, 14]
# flatten: [8, 768, 14, 14] -> [8, 768, 14x14]
# transpose: [8, 768, 14x14] -> [8, 196, 768]
x = self.proj(x).flatten(2).transpose(1, 2)
# 此处不作处理
x = self.norm(x)
return x
# Multi-Head Attention
class Attention(nn.Module):
def __init__(self,
dim, # 输入token的dim
num_heads=8,
qkv_bias=False,
qk_scale=None,
attn_drop_ratio=0.,
proj_drop_ratio=0.):
super(Attention, self).__init__()
self.num_heads = num_heads # 12
head_dim = dim // num_heads
self.scale = qk_scale or head_dim ** -0.5 # 0.125
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
self.attn_drop = nn.Dropout(attn_drop_ratio)
self.proj = nn.Linear(dim, dim)
self.proj_drop = nn.Dropout(proj_drop_ratio)
def forward(self, x):
# [batch_size, num_patches + cls_token, total_embed_dim]
# x: torch.Size([8, 197, 768])
B, N, C = x.shape
# qkv(): -> [batch_size, num_patches + 1, 3 * total_embed_dim]
# reshape: -> [batch_size, num_patches + 1, 3, num_heads, embed_dim_per_head]
# permute: -> [3, batch_size, num_heads, num_patches + 1, embed_dim_per_head]
# torch.Size([3, 8, 12, 197, 64]) 其中embed_dim_per_head = 64， num_heads * embed_dim_per_head =784
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
# [batch_size, num_heads, num_patches + 1, embed_dim_per_head]
# q,k,v: torch.Size([8, 12, 197, 64])
q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple)
# transpose: -> [batch_size, num_heads, embed_dim_per_head, num_patches + 1] -> [8, 12, 64, 197]
# @: multiply -> [batch_size, num_heads, num_patches + 1, num_patches + 1]
# torch.Size([8, 12, 197, 197])
attn = (q @ k.transpose(-2, -1)) * self.scale
attn = attn.softmax(dim=-1)
attn = self.attn_drop(attn)
# @: multiply -> [batch_size, num_heads, num_patches + 1, embed_dim_per_head] -> [8, 12, 197, 64]
# transpose: -> [batch_size, num_patches + 1, num_heads, embed_dim_per_head] -> [8, 197, 12, 64]
# reshape: -> [batch_size, num_patches + 1, total_embed_dim] -> [8, 197, 768]
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
x = self.proj(x)
x = self.proj_drop(x)
return x
# MLP层: two Linear
class Mlp(nn.Module):
"""
MLP as used in Vision Transformer, MLP-Mixer and related networks
"""
def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
super().__init__()
out_features = out_features or in_features # out_features=in_features: 768
hidden_features = hidden_features or in_features # 768*4
self.fc1 = nn.Linear(in_features, hidden_features)
self.act = act_layer()
self.fc2 = nn.Linear(hidden_features, out_features)
self.drop = nn.Dropout(drop)
# torch.Size([8, 197, 768])
def forward(self, x):
x = self.fc1(x) # torch.Size([8, 197, 3072])
x = self.act(x)
x = self.drop(x)
x = self.fc2(x) # torch.Size([8, 197, 768])
x = self.drop(x)
return x
# Transformer Encoder Block
class Block(nn.Module):
def __init__(self,
dim, # 768
num_heads,
mlp_ratio=4.,
qkv_bias=False,
qk_scale=None,
drop_ratio=0.,
attn_drop_ratio=0.,
drop_path_ratio=0.,
act_layer=nn.GELU, # 激活函数
norm_layer=nn.LayerNorm):
super(Block, self).__init__()
self.norm1 = norm_layer(dim)
# Multi-Head Attention
self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
attn_drop_ratio=attn_drop_ratio, proj_drop_ratio=drop_ratio)
# NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
self.drop_path = DropPath(drop_path_ratio) if drop_path_ratio > 0. else nn.Identity()
self.norm2 = norm_layer(dim)
# MLP
mlp_hidden_dim = int(dim * mlp_ratio) # 扩充4倍
self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop_ratio)
def forward(self, x):
x = x + self.drop_path(self.attn(self.norm1(x))) # torch.Size([8, 197, 768])
x = x + self.drop_path(self.mlp(self.norm2(x))) # torch.Size([8, 197, 768])
return x
class VisionTransformer(nn.Module):
def __init__(self,
img_size=224,
patch_size=16,
in_c=3,
num_classes=1000,
embed_dim=768,
depth=12,
num_heads=12,
mlp_ratio=4.0,
qkv_bias=True,
qk_scale=None,
representation_size=None,
distilled=False,
drop_ratio=0.,
attn_drop_ratio=0.,
drop_path_ratio=0.2,
embed_layer=PatchEmbed,
norm_layer=None,
act_layer=None):
"""
Args:
img_size (int, tuple): input image size
patch_size (int, tuple): patch size
in_c (int): number of input channels
num_classes (int): number of classes for classification head
embed_dim (int): embedding dimension
depth (int): depth of transformer
num_heads (int): number of attention heads
mlp_ratio (int): ratio of mlp hidden dim to embedding dim
qkv_bias (bool): enable bias for qkv if True
qk_scale (float): override default qk scale of head_dim ** -0.5 if set
representation_size (Optional[int]): enable and set representation layer (pre-logits) to this value if set
distilled (bool): model includes a distillation token and head as in DeiT models
drop_ratio (float): dropout rate
attn_drop_ratio (float): attention dropout rate
drop_path_ratio (float): stochastic depth rate
embed_layer (nn.Module): patch embedding layer
norm_layer: (nn.Module): normalization layer
"""
super(VisionTransformer, self).__init__()
self.num_classes = num_classes
self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models
self.num_tokens = 2 if distilled else 1 # 1 -> distilled=False
norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
act_layer = act_layer or nn.GELU
self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_c=in_c, embed_dim=embed_dim)
num_patches = self.patch_embed.num_patches
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) # torch.Size([1, 1, 768])
self.dist_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if distilled else None # 忽略，用于DeiT models
# positional embedding: torch.Size([1, 197, 768])
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
self.pos_drop = nn.Dropout(p=drop_ratio)
# 重复堆叠12次，每次的随机丢弃深度渐增
# [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
dpr = [x.item() for x in torch.linspace(0, drop_path_ratio, depth)] # stochastic depth decay rule
self.blocks = nn.Sequential(*[
Block(dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
drop_ratio=drop_ratio, attn_drop_ratio=attn_drop_ratio, drop_path_ratio=dpr[i],
norm_layer=norm_layer, act_layer=act_layer)
for i in range(depth)
])
self.norm = norm_layer(embed_dim)
# Representation layer
# representation_size: None
if representation_size and not distilled:
self.has_logits = True
self.num_features = representation_size
self.pre_logits = nn.Sequential(OrderedDict([
("fc", nn.Linear(embed_dim, representation_size)),
("act", nn.Tanh())
]))
else:
self.has_logits = False
self.pre_logits = nn.Identity() # 冻结，无操作
# Classifier head(s): Linear(in_features=768, out_features=5, bias=True)
self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
self.head_dist = None
if distilled:
self.head_dist = nn.Linear(self.embed_dim, self.num_classes) if num_classes > 0 else nn.Identity()
# Weight init
nn.init.trunc_normal_(self.pos_embed, std=0.02) # self.pos_embed: torch.Size([1, 197, 768])
if self.dist_token is not None:
nn.init.trunc_normal_(self.dist_token, std=0.02)
nn.init.trunc_normal_(self.cls_token, std=0.02) # self.cls_token: torch.Size([1, 1, 768])
self.apply(_init_vit_weights)
# x : torch.Size([B, 3, 224, 224])
def forward_features(self, x):
# [B, C, H, W] -> [B, num_patches, embed_dim]
x = self.patch_embed(x) # [B, 196, 768]
# [1, 1, 768] -> [B, 1, 768]
cls_token = self.cls_token.expand(x.shape[0], -1, -1)
if self.dist_token is None:
x = torch.cat((cls_token, x), dim=1) # [B, 197, 768]
else:
# cls_token: [B, 1, 768] concat x: [B, 196, 768] -> torch.Size([8, 197, 768])
x = torch.cat((cls_token, self.dist_token.expand(x.shape[0], -1, -1), x), dim=1)
# 添加position信息: torch.Size([8, 197, 768])
x = self.pos_drop(x + self.pos_embed)
# Transformer Encoder Block堆叠: torch.Size([8, 197, 768])
x = self.blocks(x)
x = self.norm(x)
if self.dist_token is None:
# print(x[:, 0].shape) # torch.Size([8, 768])
# self.pre_logits无操作，被冻结
return self.pre_logits(x[:, 0]) # torch.Size([8, 768])
else:
return x[:, 0], x[:, 1]
# torch.Size([8, 3, 224, 224])
def forward(self, x):
x = self.forward_features(x) # torch.Size([8, 768])
if self.head_dist is not None:
x, x_dist = self.head(x[0]), self.head_dist(x[1])
if self.training and not torch.jit.is_scripting():
# during inference, return the average of both classifier predictions
return x, x_dist
else:
return (x + x_dist) / 2
else:
x = self.head(x) # torch.Size([8, 5])
return x
def _init_vit_weights(m):
"""
ViT weight initialization
:param m: module
"""
if isinstance(m, nn.Linear):
nn.init.trunc_normal_(m.weight, std=.01)
if m.bias is not None:
nn.init.zeros_(m.bias)
elif isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode="fan_out")
if m.bias is not None:
nn.init.zeros_(m.bias)
elif isinstance(m, nn.LayerNorm):
nn.init.zeros_(m.bias)
nn.init.ones_(m.weight)
def vit_base_patch16_224_in21k(num_classes: int = 21843, has_logits: bool = True):
"""
ViT-Base model (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).
ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.
weights ported from official Google JAX impl:
https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_patch16_224_in21k-e5005f0a.pth
"""
model = VisionTransformer(img_size=224,
patch_size=16,
embed_dim=768,
depth=12,
num_heads=12,
representation_size=768 if has_logits else None,
num_classes=num_classes)
return model
def vit_base_patch32_224_in21k(num_classes: int = 21843, has_logits: bool = True):
"""
ViT-Base model (ViT-B/32) from original paper (https://arxiv.org/abs/2010.11929).
ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.
weights ported from official Google JAX impl:
https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_patch32_224_in21k-8db57226.pth
"""
model = VisionTransformer(img_size=224,
patch_size=32,
embed_dim=768,
depth=12,
num_heads=12,
representation_size=768 if has_logits else None,
num_classes=num_classes)
return model
def vit_large_patch16_224_in21k(num_classes: int = 21843, has_logits: bool = True):
"""
ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929).
ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.
weights ported from official Google JAX impl:
https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_patch16_224_in21k-606da67d.pth
"""
model = VisionTransformer(img_size=224,
patch_size=16,
embed_dim=1024,
depth=24,
num_heads=16,
representation_size=1024 if has_logits else None,
num_classes=num_classes)
return model
def vit_large_patch32_224_in21k(num_classes: int = 21843, has_logits: bool = True):
"""
ViT-Large model (ViT-L/32) from original paper (https://arxiv.org/abs/2010.11929).
ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.
weights ported from official Google JAX impl:
https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_patch32_224_in21k-9046d2e7.pth
"""
model = VisionTransformer(img_size=224,
patch_size=32,
embed_dim=1024,
depth=24,
num_heads=16,
representation_size=1024 if has_logits else None,
num_classes=num_classes)
return model
def vit_huge_patch14_224_in21k(num_classes: int = 21843, has_logits: bool = True):
"""
ViT-Huge model (ViT-H/14) from original paper (https://arxiv.org/abs/2010.11929).
ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.
NOTE: converted weights not currently available, too large for github release hosting.
"""
model = VisionTransformer(img_size=224,
patch_size=14,
embed_dim=1280,
depth=32,
num_heads=16,
representation_size=1280 if has_logits else None,
num_classes=num_classes)
return model
if __name__ == '__main__':
batch_size = 8
image_size = 224
model = vit_base_patch16_224_in21k(num_classes=5, has_logits=False)
# print(model)
image = torch.randn(batch_size, 3, image_size, image_size)
# print(image)
pred = model(image)
print(pred.shape)
pred_classes = torch.max(pred, dim=1)[1]
print(pred_classes)