参考代码地址：GitHub - google-research/vision_transformer
Transformer最开始是NLP领域提出的一篇文章，transformer因为其长依赖捕捉，出色的全局特征提取，以及关键特征相关性表征的能力，在NLP领域(比如机器翻译等)取得了巨大的成功。而Transformer中最突出的一个结构就是多头自注意力(multi-head self-attention), 通过这种token与token之间的相关性表征，来突出相互之间的重要性表征，让网络更加“关注”重要的token特征以及与其他token特征的关联。其实在CV领域很早也有了各种类型的注意力机制，比如空间注意力，通道注意力之类的。而VIT是首次，保留了一个完整的transformer的结构，应用到CV领域。
文章中的实验表示VIT在足够的预训练的后，在分类任务上的表现要好于CNN网络结构的。接下来介绍一下VIT的具体方法。
1、 原理简介：
我们来看一下这张图(从原文中截取出来的)：
这个是模型的一个overview，看一下这段文字描述，大概的意思：
1. 我们把一张图，分成固定大小块(patch), 然后对这些patch进行线性化的embedding，然后加上位置embedding。然后把这个分完块后的序列性的向量 送到原始Transformer的encoder中。
2. 为了实现分类，我们由在这个序列向量中加入一个"分类token".
主要原理到此为止，讲完了。是不是非常简单，确实很简单。我们再来看图，再进一步理解。虚线左半部分是vit的示意图。右边部分是原始transformer的encoder的结构示意图。
1. 左边从下往上，先图像分成patch，再把二维排列patch拉平(2D flatten to 1D), 排列成一维的序列。
2. 然后把这些对这些patch进行线性化embedding， 再加上位置embedding。得到如绿色框框出来的那样，有数字的表示位置embedding,无数字的表示patch embedding。然后最左边红色框框出来的是用来做分类的embedding。
3. 然后把这个序列向量喂给原始transformer的encoder。而原始transformer的encoder如虚线右边图所示(原始transformer这里就不详细展开了，图中也非常清楚，Norm层 + multi-head self-attention接残差，再接Norm层 + MLP接残差，transformer原理之后会单独出一篇文章介绍)。下面的直接进入代码部分就很清楚了。
4. transformer encoder结构出来的结果，再送到一个分类head(这个分类head是一个MLP结构)中，最终出最后的分类结果。
2、代码详解
读者可以先参考一下官方给的代码，这里笔者手写了一份pytorch版本，提供参考，代码思路就是按照上面原理简介来的。
2.1 整体结构：
class VisionTransformer(nn.Module):
def __init__(self, in_channel=3, embed_dim=64, patch_size=(4, 4),
input_res=(224, 224), num_classes=10, num_layers=2,
num_head=6, head_dim=64, hidden_dim=128, out_dim=64, drop_out=0.5):
super(VisionTransformer, self).__init__()
num_patch = (input_res[0] // patch_size[0]) * (input_res[1] // patch_size[1])
self.linear_embedding = Embedding(in_channel, embed_dim,
num_patch, patch_size)
self.encoders = nn.Sequential(*[
EncoderBlock(
embed_dim, out_dim, num_head=num_head,
head_dim=head_dim, hidden_dim=hidden_dim,
attn_drop_rate=drop_out, drop_rate=drop_out)
for _ in range(num_layers)])
self.head = MLP(input_dim=out_dim, hidden_dim=(out_dim // 2), out_dim=num_classes)
def forward(self, x):
x = self.linear_embedding(x)
x = self.encoders(x)
x = x[:, 0, :] # 取出用于做分类的那个向量
logits = self.head(x)
return logits看forward，首先进行embedding, 这里需要定义一个embedding的方法。然后把embedding的输出传入transformer的encoder中，这里的encoder可以串联多个，在self.encoders定义中用num_layers这个参数控制个数，这里需要定义一个EncoderBlock。之后，encoder的输出，我们把用于做分类的那个向量拿出来，给到分类的head，最后得到logits。代码流程和原理简介中的图一模一样。接下来详细看embedding和EncoderBlock的定义。
2.2 Embedding
class Embedding(nn.Module):
def __init__(self, input_dim, embed_dim, num_patch, patch_size):
super(Embedding, self).__init__()
self.num_patch = num_patch
self.patch_size = patch_size
# 这里直接采用卷积对输入进行线性embedding，
self.proj_embedding = nn.Conv2d(in_channels=input_dim,
out_channels=embed_dim,
kernel_size=patch_size,
stride=patch_size)
# 这里采用最简单的参数进行position embedding, 也可以是其他的方式
self.position_embedding = nn.Parameter(torch.ones(1, num_patch, embed_dim))
# 示意图中红色框框的部分，用了一个单独的embedding做分类
self.class_embedding = nn.Parameter(torch.zeros(1, 1, embed_dim))
def forward(self, x):
b, _, H, W = x.shape
patch_n = (H // self.patch_size[0]) * (W // self.patch_size[1])
assert patch_n == self.num_patch, "input patch number is not equal to defined patch number"
x = self.proj_embedding(x) # embedding for input
x = x.view(b, -1, patch_n).permute(0, 2, 1)
x = x + self.position_embedding # add position embedding
x = torch.cat((self.class_embedding, x), dim=1) # concat class embedding at dim of number patch
return x2.3 EncoderBlock
class EncoderBlock(nn.Module):
def __init__(self, input_dim, out_dim, num_head=6, head_dim=64, hidden_dim=64,
attn_drop_rate=0.5, drop_rate=0.5
):
super(EncoderBlock, self).__init__()
self.attentions = MultiHeadSelfAttention(input_dim, num_head, head_dim, attn_drop_rate)
self.mlp = MLP(input_dim, hidden_dim, out_dim, drop_rate)
def forward(self, x):
y = self.attentions(x)
x = x + y
y = self.mlp(x)
y = x + y
return y2.3.1 Mult-Head Self-Attention
class MultiHeadSelfAttention(nn.Module):
def __init__(self, input_dim, num_head=6, head_dim=64, drop_rate=0.5):
super(MultiHeadSelfAttention, self).__init__()
self.num_head = num_head
self.norm = nn.LayerNorm(input_dim)
self.qs = nn.ModuleList([nn.Linear(input_dim, head_dim) for _ in range(num_head)])
self.ks = nn.ModuleList([nn.Linear(input_dim, head_dim) for _ in range(num_head)])
self.vs = nn.ModuleList([nn.Linear(input_dim, head_dim) for _ in range(num_head)])
self.projection = nn.Sequential(
nn.Linear(head_dim * num_head, input_dim),
nn.GELU(),
nn.Dropout(drop_rate)
)
self.scale = head_dim ** -0.5
def forward(self, x): # x shape: (B, N, C) B: batch_size, N: number patch, C: embedding dim
attn_outs = []
y = self.norm(x)
for i in range(self.num_head):
q = self.qs[i](y)
k = self.ks[i](y)
v = self.vs[i](y)
attn = torch.matmul(q, k.transpose(-2, -1))
attn = attn.softmax(dim=-1)
attn_out = torch.matmul(attn, v)
attn_outs.append(attn_out)
y = torch.cat(attn_outs, dim=-1)
y = self.projection(y)
return y2.3.2 MLP
class MLP(nn.Module):
def __init__(self, input_dim, hidden_dim, out_dim, drop_rate=0.5):
super(MLP, self).__init__()
self.mlp = nn.Sequential(
nn.LayerNorm(input_dim),
nn.Linear(input_dim, hidden_dim),
nn.GELU(),
nn.Dropout(drop_rate),
nn.Linear(hidden_dim, out_dim),
nn.GELU(),
nn.Dropout(drop_rate),
)
def forward(self, x):
y = self.mlp(x)
return y至此，代码部分全部写完了，看上去也没有那么复杂哈。当然这个也只是笔者参考论文的描述即兴写的，运行起来肯定是没问题的，大家可以参考一下。但是详细的一些参数设置或者是embedding方法可以是不同的。而且最后这个分类的head，也可以的不同的。而且论文中用的是一个单独的token embedding的分类向量的方式用来做head的输入，当然也可以不同单独用一个token embedding做分类向量，也可以用一些adaptavgpool的方式也是也可以的。包括，如果不是做分类任务，如果是其他的任务，这个head也是可以灵活替换的。
论文中的实验以及其他的内容这里就不再介绍了，大家自行看论文了解哈。最后，觉得写的不错的同学可以点赞收藏加个关注，谢谢各位同学，咱们下一篇文章见！
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
