注：
Norm_layer 参数我们传入的是 nn.LayerNorm.
我们查看下layerNorm的信息：
self.norm = norm_layer(embed_dim) print(self.norm)
LayerNorm((768,), eps=1e-06, elementwise_affine=True)
Key Code:
training head.bias
0%| | 0/368 [00:00<?, ?it/s]tensor([ 6.4645, -4.7138, 4.6655, -5.7235, -2.1071, -3.0778, -3.2277, 1.0825,
-3.1296, 6.2348], device='cuda:0')
tensor(-0.0354, device='cuda:0') tensor(6.8029, device='cuda:0')
tensor([ 0.1896, -0.1529, 0.1320, -0.1904, -0.0908, -0.1010, -0.1128, 0.0217,
-0.1351, 0.2111], device='cuda:0')
z= tensor([ 0.1895, -0.1528, 0.1319, -0.1902, -0.0907, -0.1009, -0.1128, 0.0217,
-0.1350, 0.2109], device='cuda:0')
torch.Size([8, 197, 768])
Total Code:
"""
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
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: # if there is no throwing or no 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) # torch.rand() [0,1)
random_tensor.floor_() # binarize # random_tensor.floor_() 会将每个浮点数向下取整
output = x.div(keep_prob) * random_tensor
return output
# x = torch.tensor([[1.0, 2.0],
# [3.0, 4.0],
# [5.0, 6.0],
# [7.0, 8.0]])
# keep_prob = 1 - drop_prob
# keep_prob = 1 - 0.5
# keep_prob = 0.5
# shape = (x.shape[0],) + (1,) * (x.ndim - 1)
# shape = (4,) + (1,) * (2 - 1)
# shape = (4, 1)
#
# random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
# random_tensor = 0.5 + torch.rand((4, 1))
#
# # 生成的 random_tensor 可能类似于
# random_tensor = torch.tensor([[0.8],
# [0.3],
# [0.7],
# [0.1]])
# random_tensor.floor_()
# random_tensor = torch.tensor([[1.0],
# [0.0],
# [1.0],
# [0.0]])
# output = x.div(keep_prob) * random_tensor
# output = x.div(0.5) * random_tensor
#
# # 计算每个元素
# output = torch.tensor([[1.0 / 0.5, 2.0 / 0.5],
# [3.0 / 0.5, 4.0 / 0.5],
# [5.0 / 0.5, 6.0 / 0.5],
# [7.0 / 0.5, 8.0 / 0.5]]) * random_tensor
#
# output = torch.tensor([[2.0, 4.0],
# [6.0, 8.0],
# [10.0, 12.0],
# [14.0, 16.0]]) * random_tensor
#
# # 计算结果
# output = torch.tensor([[2.0, 4.0],
# [0.0, 0.0],
# [10.0, 12.0],
# [0.0, 0.0]])
# 在这个例子中，每个样本的路径被随机丢弃。
# 对于丢弃的路径，输出张量中的对应元素被设置为0，而未丢弃的路径的值按比例放大，以保持整体的期望值不变。
#
# 这种方法在深度学习中的作用类似于 Dropout，但它是对整个路径（例如残差块）进行随机丢弃，而不是对单个神经元。
# 这种技术能够提高模型的鲁棒性和泛化能力。
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) # (224,224)
patch_size = (patch_size, patch_size) # (16,16)
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] # 14*14=196
self.proj = nn.Conv2d(in_c, embed_dim, kernel_size=patch_size, stride=patch_size) # 3,768,16,16
self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
def forward(self, x):
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]})."
# flatten: [B, C, H, W] -> [B, C, HW]
# transpose: [B, C, HW] -> [B, HW, C]
x = self.proj(x).flatten(2).transpose(1, 2)
x = self.norm(x)
return x
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 # 64 = 768 / 12
self.scale = qk_scale or head_dim ** -0.5 # 这是计算 64 的倒数的平方根
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) # on the last dimension of the input tensor
self.attn_drop = nn.Dropout(attn_drop_ratio)
# 在训练模式下（model.train()），Dropout 层以 50% 的概率随机将一些神经元的输出置为零。
# 在评估模式下（model.eval()），Dropout 层不会丢弃任何神经元的输出，但会将所有神经元的输出乘以 0.5（因为丢弃概率为 0.5）。
self.proj = nn.Linear(dim, dim)
self.proj_drop = nn.Dropout(proj_drop_ratio)
def forward(self, x):
# [batch_size, num_patches + 1, total_embed_dim]
B, N, C = x.shape
# print(x.shape) # torch.Size([8, 197, 768])
# print(self.qkv(x).shape) # torch.Size([8, 197, 2304])
# print(self.num_heads) # 12
# print(self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).shape)# torch.Size([8, 197, 3, 12, 64])
# input()
# qkv(): -> [batch_size, num_patches + 1, 3 * total_embed_dim]
# reshape: -> [batch_size, num_patches + 1, 3, num_heads, embed_dim_per_head] # # torch.Size([8, 197, 3, 12, 64])
# permute: -> [3, batch_size, num_heads, num_patches + 1, embed_dim_per_head] # # torch.Size([3, 8, 12, 197, 64])
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 = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) # torch.Size([8, 12, 197, 64])
# print(q.shape) # torch.Size([8, 12, 197, 64])
# print(q.transpose(-2,-1).shape) # torch.Size([8, 12, 64, 197])
# print((q.permute(0,1,3,2)).shape) # torch.Size([8, 12, 64, 197])
#
# # print(q.permute(-2,-1)) # error: i.e. input.dim() = 4 is not equal to len(dims) = 2
#
# print("***")
#
# print((q @ k.transpose(-2, -1)) .shape) # torch.Size([8, 12, 197, 197])
#
# input()
# print(64 ** -0.5) # 这是计算 64 的倒数的平方根。
# input()
# transpose: -> [batch_size, num_heads, embed_dim_per_head, num_patches + 1]
# @: multiply -> [batch_size, num_heads, num_patches + 1, num_patches + 1]
attn = (q @ k.transpose(-2, -1)) * self.scale # torch.Size([8, 12, 197, 197])
attn = attn.softmax(dim=-1) # softmax 是一种常用的激活函数，通常用于多分类任务的输出层。它将输入的向量转换为概率分布，即输出的每个元素都在 0 到 1 之间，并且这些元素的总和为 1
attn = self.attn_drop(attn)
# @: multiply -> [batch_size, num_heads, num_patches + 1, embed_dim_per_head] # torch.Size([8, 12, 197, 64])
# transpose: -> [batch_size, num_patches + 1, num_heads, embed_dim_per_head] # torch.Size([8, 197, 12, 64])
# reshape: -> [batch_size, num_patches + 1, total_embed_dim] # torch.Size([8, 197, 768])
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
x = self.proj(x) # torch.Size([8, 197, 768])
x = self.proj_drop(x)
return x
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
hidden_features = hidden_features or in_features
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)
def forward(self, x):
x = self.fc1(x)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return x
class Block(nn.Module):
def __init__(self,
dim,
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)
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_hidden_dim = int(dim * mlp_ratio)
self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop_ratio)
def forward(self, x):
# torch.Size([8, 197, 768]) # Batch=8 对每个样本的所有通道计算均值和标准差
x = x + self.drop_path(self.attn(self.norm1(x))) # Attention input(x): torch.Size([8,197,768]) output: torch.Size([8,197,768])
x = x + self.drop_path(self.mlp(self.norm2(x)))
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., 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
norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
# partial 是 Python functools 模块中的一个函数，它可以用来固定函数的一些参数，返回一个新的函数。
# 这个新的函数可以像原函数一样被调用，但已经有部分参数被固定住了。
# partial(nn.LayerNorm, eps=1e-6) 的意思是创建一个新的 LayerNorm 类，其中 eps 参数被固定为 1e-6。
# 这个新的类可以像 LayerNorm 类一样使用，但不需要每次都指定 eps 参数，因为它已经被固定了。
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))
# 定义一个可训练的分类标记（classification token，简称 cls_token）。这个标记用于表示整个输入序列或图像的全局信息，通常用于分类任务。
self.dist_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if distilled else None
# 定义了一个形状为 (1, 1, embed_dim) 的全零张量，并将其转换为可训练的参数（如果 distilled 为 True）。
# 初始化为全零向量，并将其转换为可训练参数。
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
self.pos_drop = nn.Dropout(p=drop_ratio)
# 这行代码是用于实现随机深度（Stochastic Depth）技术的一部分。
# 在Vision Transformer (ViT)模型中，随机深度是一种正则化方法，
# 它通过在训练期间随机丢弃某些层以提高模型的泛化能力。
# 具体来说，这段代码生成了一系列在[0, drop_path_ratio]之间均匀分布的值，用于控制每个Transformer层的drop path概率。
# tensor([ 0.0000, 454.5454, 909.0909, 1363.6364, 1818.1818, 2272.7273,
# 2727.2727, 3181.8184, 3636.3638, 4090.9092, 4545.4546, 5000.0000])
# print(torch.linspace(0, 5000, depth))
# [0.0, 454.5454406738281, 909.0908813476562, 1363.6363525390625, 1818.1817626953125, 2272.727294921875,
# 2727.272705078125, 3181.818359375, 3636.36376953125, 4090.9091796875, 4545.45458984375, 5000.0]
# print([x.item() for x in torch.linspace(0, 5000, depth)])
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)
])
# LayerNorm在Transformer模型中的应用是对每个样本的特征维度进行归一化，而不是对所有token或patch的维度进行归一化。
# 这种设计有助于模型更好地捕捉每个token或patch的特征，提高训练的稳定性和收敛速度。
# LayerNorm((768,), eps=1e-06, elementwise_affine=True)
self.norm = norm_layer(embed_dim)
# Representation layer
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)
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
# 良好的权重初始化是深度学习模型训练中至关重要的一步。通过上述初始化策略，可以：
#
# 减少训练初期的随机噪声，提高模型收敛速度。
# 防止梯度消失或爆炸，保持训练过程的稳定性。
# 确保特定类型层（如位置嵌入、分类标记和蒸馏标记）的初始状态合理。
# 使用截断正态分布（trunc_normal_），标准差为 0.02。位置嵌入的初始化对 Transformer 模型中的位置信息编码非常重要。
nn.init.trunc_normal_(self.pos_embed, std=0.02)
if self.dist_token is not None:
# 同样使用截断正态分布（trunc_normal_），标准差为 0.02。蒸馏标记的初始化确保其在训练开始时不会引入过多噪声。
nn.init.trunc_normal_(self.dist_token, std=0.02)
# 同样使用截断正态分布（trunc_normal_），标准差为 0.02。分类标记用于表示整个输入序列的全局信息，因此初始化也非常重要。
nn.init.trunc_normal_(self.cls_token, std=0.02)
# self.apply(_init_vit_weights)：递归地将 _init_vit_weights 函数应用于模型的所有子模块。
# 这一步确保模型中所有线性层、卷积层和归一化层都按照定义的方式进行初始化。
self.apply(_init_vit_weights)
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:
x = torch.cat((cls_token, self.dist_token.expand(x.shape[0], -1, -1), x), dim=1)
x = self.pos_drop(x + self.pos_embed)
x = self.blocks(x)
print(x[0,0,0:10])
gamma = self.norm.weight
beta = self.norm.bias
# print(gamma,beta)
y=x
mean=torch.mean(y[0,0,:])
std=torch.std(y[0,0,:])
z=(y[0,0,:]-mean)/std
z = z * gamma + beta
print(mean,std)
x = self.norm(x)
print(x[0,0,0:10])
print("z=",z[0:10])
print(x.shape)
input()
if self.dist_token is None:
return self.pre_logits(x[:, 0])
else:
return x[:, 0], x[:, 1]
# x[:, 0] 是每个样本的分类标记
# x[:, 1] 是每个样本的蒸馏标记（dist_token）
def forward(self, x):
x = self.forward_features(x)
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)
return x
def _init_vit_weights(m):
"""
ViT weight initialization
:param m: module
"""
# nn.Linear 层：
#
# 权重初始化：使用截断正态分布（trunc_normal_），标准差为 0.01。这种初始化方法通过截断分布的尾部减少了极端值，从而有助于模型更快地收敛。
# 偏置初始化：将偏置初始化为零（zeros_）。
# nn.Conv2d 层：
#
# 权重初始化：使用 Kaiming 正态初始化（kaiming_normal_），模式为 fan_out。这种方法通常用于 ReLU 激活函数，有助于保持前向传播过程中方差的一致性。
# 偏置初始化：将偏置初始化为零（zeros_）。
# nn.LayerNorm 层：
#
# 权重初始化：将权重初始化为一（ones_）。
# 偏置初始化：将偏置初始化为零（zeros_）。
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(num_classes: int = 1000):
"""
ViT-Base model (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).
ImageNet-1k weights @ 224x224, source https://github.com/google-research/vision_transformer.
weights ported from official Google JAX impl:
链接: https://pan.baidu.com/s/1zqb08naP0RPqqfSXfkB2EA 密码: eu9f
"""
model = VisionTransformer(img_size=224,
patch_size=16,
embed_dim=768,
depth=12,
num_heads=12,
representation_size=None,
num_classes=num_classes)
return model
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(num_classes: int = 1000):
"""
ViT-Base model (ViT-B/32) from original paper (https://arxiv.org/abs/2010.11929).
ImageNet-1k weights @ 224x224, source https://github.com/google-research/vision_transformer.
weights ported from official Google JAX impl:
链接: https://pan.baidu.com/s/1hCv0U8pQomwAtHBYc4hmZg 密码: s5hl
"""
model = VisionTransformer(img_size=224,
patch_size=32,
embed_dim=768,
depth=12,
num_heads=12,
representation_size=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(num_classes: int = 1000):
"""
ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929).
ImageNet-1k weights @ 224x224, source https://github.com/google-research/vision_transformer.
weights ported from official Google JAX impl:
链接: https://pan.baidu.com/s/1cxBgZJJ6qUWPSBNcE4TdRQ 密码: qqt8
"""
model = VisionTransformer(img_size=224,
patch_size=16,
embed_dim=1024,
depth=24,
num_heads=16,
representation_size=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
4万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
