ViT论文Pytorch代码解读
ViT论文代码实现
论文地址：https://arxiv.org/abs/2010.11929
Pytorch代码地址：https://github.com/lucidrains/vit-pytorch
ViT结构图
调用代码
import torch
from vit_pytorch import ViT
def test():
v = ViT(
image_size = 256,
patch_size = 32,
num_classes = 1000,
dim = 1024,
depth = 6,
heads = 16,
mlp_dim = 2048,
dropout = 0.1,
emb_dropout = 0.1
)
img = torch.randn(1, 3, 256, 256)
preds = v(img)
print(preds.shape)
assert preds.shape == (1, 1000), 'correct logits outputted'
if __name__ == '__main__':
test()
ViT结构
class ViT(nn.Module):
def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool='cls', channels=3,
dim_head=64, dropout=0., emb_dropout=0.):
super().__init__()
# 将image_size和patch_size都转换为(height, width)形式
image_height, image_width = pair(image_size)
patch_height, patch_width = pair(patch_size)
# 检查图像尺寸是否可以被patch尺寸整除
assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
# 计算图像中的patch数量
num_patches = (image_height // patch_height) * (image_width // patch_width)
# 计算每个patch的维度（即每个patch的元素数量）
patch_dim = channels * patch_height * patch_width
# 确保池化方式是'cls'或'mean'
assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
# 将图像转换为patch嵌入的操作
self.to_patch_embedding = nn.Sequential(
Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=patch_height, p2=patch_width), # 图像切分重排，后文有注释
# 注：此时的维度为[b, h*w/p1/p2, p1*p2*c]:[批处理尺寸、图像中patch的数、每个patch的元素数量]
nn.LayerNorm(patch_dim), # 对patch进行层归一化
nn.Linear(patch_dim, dim), # 使用线性层将patch的维度从patch_dim转化为dim
nn.LayerNorm(dim), # 对结果进行层归一化
)
self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim)) # 初始化位置嵌入
self.cls_token = nn.Parameter(torch.randn(1, 1, dim)) # 初始化CLS token（用于分类任务的特殊token）
self.dropout = nn.Dropout(emb_dropout)
self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout) # 定义Transformer模块
self.pool = pool # 设置池化方式（'cls'或'mean'）
self.to_latent = nn.Identity() # 设置一个恒等映射（在此实现中不改变数据，但可以在子类或其他变种中进行修改）
self.mlp_head = nn.Linear(dim, num_classes) # 定义MLP头部，用于最终的分类
def forward(self, img):
x = self.to_patch_embedding(img) # 第一步，将图片切分为若干小块
# 此时维度为：[b, h*w/p1/p2, dim]
b, n, _ = x.shape
# 第二步，设置位置编码
cls_tokens = repeat(self.cls_token, '1 1 d -> b 1 d', b=b) # 将cls_token复制b个
# (为每个输入图像复制一个CLS token，使输入批次中的每张图像都有一个相应的CLS token)
x = torch.cat((cls_tokens, x), dim=1) # 将CLS token与patch嵌入合并; cat之后，原来的维度[1,64,1024],就变成了[1,65,1024]
x += self.pos_embedding[:, :(n + 1)] # 原数据和位置编码直接进行相加操作，即完成结构图中的【Patch + Position Embedding】操作
x = self.dropout(x)
# 第三步，Transformer的Encoder结构
x = self.transformer(x)
x = x.mean(dim=1) if self.pool == 'mean' else x[:, 0] # 根据所选的池化方式进行池化
x = self.to_latent(x) # 将数据传递给恒等映射
return self.mlp_head(x) # 使用MLP头部进行分类
Rearrange解释：
y = x.transpose(0, 2, 3, 1)
可以写成：y = rearrange(x, ‘b c h w -> b h w c’)
关于pos_embedding和cls_token的逻辑讲解：
如图所示，红色框框出的部分。
图像被切分为多个小块之后，经过self.to_patch_embedding中的Rearrange，原本的[b,c,h,w]维度变为[b, h*w/p1/p2, p1*p2*c]。
再经过线性层nn.Linear(patch_dim, dim)，维度变为[b, h*w/p1/p2, dim]。
输出结果即为上图中黄色框标出的部分的粉色条（不包括紫色条，是因为此处还没进行Position Embedding操作）。
继续往下走，进行torch.cat((cls_tokens, x), dim=1)，此时将x与cls_tokens进行concat操作，得到红色框框出的所有粉色条（在原本的基础上增加了带*号的粉色条）。
记下来的x += self.pos_embedding[:, :(n + 1)]操作就是将x与pos_embedding直接进行相加，用图表示出来就是上图中整个红色框框出的部分了（紫色条就是传说中的pos_embedding）。
举一个有数字的例子：
原本输入图像维度为[1, 3, 256, 256]，dim设置为1023，经过self.to_patch_embedding后维度变为：[1,64,1024]，cls_tokens的维度为：[1,1,1024]，经过concat操作后，x的维度变为[1,65,1024]，然后经过pos_embedding加操作后，维度依然是[1,65,1024]，因为在设置变量pos_embedding时的维度就是torch.randn(1, num_patches + 1, dim)。
~这个解释应该够清晰了吧！~
Transformer Encoder结构
# 定义前馈神经网络
class FeedForward(nn.Module):
def __init__(self, dim, hidden_dim, dropout=0.):
super().__init__()
self.net = nn.Sequential(
# Vit_base: dim=768,hidden_dim=3072
nn.LayerNorm(dim),
nn.Linear(dim, hidden_dim), # 将输入从dim维映射到hidden_dim维
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(hidden_dim, dim), # 将隐藏状态从hidden_dim维映射回到dim维
nn.Dropout(dropout)
)
def forward(self, x):
return self.net(x)
class Attention(nn.Module):
def __init__(self, dim, heads=8, dim_head=64, dropout=0.):
super().__init__()
inner_dim = dim_head * heads # 64*8=512 # 计算内部维度
project_out = not (heads == 1 and dim_head == dim) # 判断是否需要投影输出，投影输出就是是否需要经过线性层
# 如果只有一个attention头并且其维度与输入相同则不需要投影输出，否则需要。
self.heads = heads
self.scale = dim_head ** -0.5 # 缩放因子，通常是头维度的平方根的倒数
self.norm = nn.LayerNorm(dim)
self.attend = nn.Softmax(dim=-1) # softmax函数用于最后一个维度，计算注意力权重
self.dropout = nn.Dropout(dropout)
self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False) # 一个线性层生成Q, K, V
# 判断是否需要投影输出
self.to_out = nn.Sequential(
nn.Linear(inner_dim, dim),
nn.Dropout(dropout)
) if project_out else nn.Identity()
def forward(self, x):
x = self.norm(x)
qkv = self.to_qkv(x).chunk(3, dim=-1) # 用线性层生成QKV，并在最后一个维度上分块；相当于写3遍nn.Linear
q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.heads), qkv)
# 将[batch_size, sequence_length, heads_dimension] 转换为 [batch_size, number_of_heads, sequence_length, dimension_per_head]
dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale # 计算Q和K的点乘，然后进行缩放
# q: [batch_size, number_of_heads, sequence_length, dimension_per_head]
# k转置后：[batch_size, number_of_heads, sequence_length, dimension_per_head] -> [batch_size, number_of_heads, dimension_per_head, sequence_length]
# q和k点乘后：[batch_size, number_of_heads, sequence_length, sequence_length]
attn = self.attend(dots) # 使用softmax函数获取注意力权重
attn = self.dropout(attn)
# 使用注意力权重对V进行加权
out = torch.matmul(attn, v)
out = rearrange(out, 'b h n d -> b n (h d)') # 使用rearrange函数重新组织输出的维度
return self.to_out(out) # 投影输出（如果需要）
class Transformer(nn.Module):
def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.):
super().__init__()
self.norm = nn.LayerNorm(dim)
self.layers = nn.ModuleList([])
for _ in range(depth): # depth设置为几层，就重复几次
self.layers.append(nn.ModuleList([
Attention(dim, heads=heads, dim_head=dim_head, dropout=dropout),
FeedForward(dim, mlp_dim, dropout=dropout)
]))
def forward(self, x):
for attn, ff in self.layers: # 残差
x = attn(x) + x
x = ff(x) + x
return self.norm(x)
如上就是ViT的整体结构了。
附：完整代码
import torch
from torch import nn
from einops import rearrange, repeat
from einops.layers.torch import Rearrange
# helpers
def pair(t):
return t if isinstance(t, tuple) else (t, t)
# classes
class FeedForward(nn.Module):
def __init__(self, dim, hidden_dim, dropout=0.):
super().__init__()
self.net = nn.Sequential(
# Vit_base: dim=768,hidden_dim=3072
nn.LayerNorm(dim),
nn.Linear(dim, hidden_dim),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(hidden_dim, dim),
nn.Dropout(dropout)
)
def forward(self, x):
return self.net(x)
class Attention(nn.Module):
def __init__(self, dim, heads=8, dim_head=64, dropout=0.):
super().__init__()
inner_dim = dim_head * heads # 64*8=512
project_out = not (heads == 1 and dim_head == dim)
self.heads = heads
self.scale = dim_head ** -0.5
self.norm = nn.LayerNorm(dim)
self.attend = nn.Softmax(dim=-1)
self.dropout = nn.Dropout(dropout)
self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)
self.to_out = nn.Sequential(
nn.Linear(inner_dim, dim),
nn.Dropout(dropout)
) if project_out else nn.Identity()
def forward(self, x):
x = self.norm(x)
qkv = self.to_qkv(x).chunk(3, dim=-1) # 相当于写3遍nn.Linear
q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.heads), qkv)
# 将[batch_size, sequence_length, heads_dimension] 转换为 [batch_size, number_of_heads, sequence_length, dimension_per_head]
dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
# q: [batch_size, number_of_heads, sequence_length, dimension_per_head]
# k转置后：[batch_size, number_of_heads, sequence_length, dimension_per_head] -> [batch_size, number_of_heads, dimension_per_head, sequence_length]
# q和k点乘后：[batch_size, number_of_heads, sequence_length, sequence_length]
attn = self.attend(dots)
attn = self.dropout(attn)
out = torch.matmul(attn, v)
out = rearrange(out, 'b h n d -> b n (h d)')
return self.to_out(out)
class Transformer(nn.Module):
def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.):
super().__init__()
self.norm = nn.LayerNorm(dim)
self.layers = nn.ModuleList([])
for _ in range(depth):
self.layers.append(nn.ModuleList([
Attention(dim, heads=heads, dim_head=dim_head, dropout=dropout),
FeedForward(dim, mlp_dim, dropout=dropout)
]))
def forward(self, x):
for attn, ff in self.layers:
x = attn(x) + x
x = ff(x) + x
return self.norm(x)
class ViT(nn.Module):
def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool='cls', channels=3,
dim_head=64, dropout=0., emb_dropout=0.):
super().__init__()
image_height, image_width = pair(image_size)
patch_height, patch_width = pair(patch_size)
assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
num_patches = (image_height // patch_height) * (image_width // patch_width)
patch_dim = channels * patch_height * patch_width
assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
self.to_patch_embedding = nn.Sequential(
Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=patch_height, p2=patch_width), # 图像切分重排
nn.LayerNorm(patch_dim),
nn.Linear(patch_dim, dim),
nn.LayerNorm(dim),
)
# Rearrange解释：
# y = x.transpose(0, 2, 3, 1)
# 可以写成：y = rearrange(x, 'b c h w -> b h w c')
self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
self.dropout = nn.Dropout(emb_dropout)
self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
self.pool = pool
self.to_latent = nn.Identity()
self.mlp_head = nn.Linear(dim, num_classes)
def forward(self, img):
x = self.to_patch_embedding(img)
b, n, _ = x.shape
cls_tokens = repeat(self.cls_token, '1 1 d -> b 1 d', b=b) # 数字编码，将cls_token复制b个
x = torch.cat((cls_tokens, x), dim=1) # cat之后，原来的维度[1,64,1024],就变成了[1,65,1024]
x += self.pos_embedding[:, :(n + 1)]
x = self.dropout(x)
x = self.transformer(x)
x = x.mean(dim=1) if self.pool == 'mean' else x[:, 0]
x = self.to_latent(x)
return self.mlp_head(x)
附：训练代码
model = ViT(
dim=128,
image_size=224,
patch_size=32,
num_classes=2,
transformer=efficient_transformer,
channels=3,
).to(device)
# loss function
criterion = nn.CrossEntropyLoss()
# optimizer
optimizer = optim.Adam(model.parameters(), lr=lr)
# scheduler
scheduler = StepLR(optimizer, step_size=1, gamma=gamma)
for epoch in range(epochs):
epoch_loss = 0
epoch_accuracy = 0
for data, label in tqdm(train_loader):
data = data.to(device)
label = label.to(device)
output = model(data)
loss = criterion(output, label)
optimizer.zero_grad()
loss.backward()
optimizer.step()
acc = (output.argmax(dim=1) == label).float().mean()
epoch_accuracy += acc / len(train_loader)
epoch_loss += loss / len(train_loader)
with torch.no_grad():
epoch_val_accuracy = 0
epoch_val_loss = 0
for data, label in valid_loader:
data = data.to(device)
label = label.to(device)
val_output = model(data)
val_loss = criterion(val_output, label)
acc = (val_output.argmax(dim=1) == label).float().mean()
epoch_val_accuracy += acc / len(valid_loader)
epoch_val_loss += val_loss / len(valid_loader)
print(
f"Epoch : {epoch+1} - loss : {epoch_loss:.4f} - acc: {epoch_accuracy:.4f} - val_loss : {epoch_val_loss:.4f} - val_acc: {epoch_val_accuracy:.4f}\n"
)