vit-transformer模型结构及源码解读

vit简介
vit模型是transformer在图像分类领域的首次成功尝试；
但是其需要基于大量数据去预训练模型；除了训练难度，现有的 Visual Transformer 参数量和计算量多大，比如 ViT 需要 18B FLOPs 在 ImageNet 达到 78% 左右 Top1，但是 CNN 模型如 GhostNet 只需 600M FLOPs 可以达到 79% 以上 Top1。

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
#PreNorm是对层进行归一化
class PreNorm(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.fn = fn
    def forward(self, x, **kwargs):
        return self.fn(self.norm(x), **kwargs)
#FeedForward就是两层线性变换
class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout = 0.):
        super().__init__()
        self.net = nn.Sequential(
            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)
#attention的输入和输出维度相同[1,num_patches+1,128]-->[num_patches+1,128],其目的是赋予不同patch不同的权重；
#给予不同的注意力
class Attention(nn.Module):
    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
        super().__init__()
        inner_dim = dim_head *  heads
        project_out = not (heads == 1 and dim_head == dim)

        self.heads = heads
        self.scale = dim_head ** -0.5

        self.attend = nn.Softmax(dim = -1)
        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):
        qkv = self.to_qkv(x).chunk(3, dim = -1)
        #获得三个维度相同的向量q,k,v,然后q,k相乘获得权重，乘以scale,再经过softmax之后，乘到v上
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)

        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale

        attn = self.attend(dots)

        out = torch.matmul(attn, v)
        out = rearrange(out, 'b h n d -> b n (h d)')
        return self.to_out(out)
#Transformer就是将降维后的patches叠加上不同的系数（注意力机制），再加上两层线性传输
class Transformer(nn.Module):
    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):
        super().__init__()
        self.layers = nn.ModuleList([])
        for _ in range(depth):
            self.layers.append(nn.ModuleList([
                PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout)),
                PreNorm(dim, 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 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)'
        #这里是对块进行编码，将patch_height*patch_width的大小输出维度变成隐层dim
        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.Linear(patch_dim, dim),
        )
        #这里对每一个块进行位置编码(包含cls_tokens),位置编码的长度为128，会以相加的方式叠加到原始输入上
        #关于位置编码
        #1.固定正弦编码 - 无学习参数
        #2.绝对位置编码 - 一维学习编码
        #3.轴向位置编码 - 二维学习编码
        # 大多数 NLP 模型（和 GPT）只使用 2。是的，新的视觉 SOTA 与 GPT 的架构相同，只有细微的差别。条条大路通罗马。
        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
        #关于cls_token，原先在NLP中，这个主要用于生成整个句子的分类，但是看过一些解释，去掉这个，用平均值或者加权值代表整个句子的分类也是一样的
        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.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, num_classes)
        )

    def forward(self, img):
        #对图像进行分块和降维编码
        x = self.to_patch_embedding(img)
        b, n, _ = x.shape
        #加上一个分类维度[1,1,128]叠加到输入上面
        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)
        #对num_patches+1这个维度求均值，x的维度由[1,num_patches+1,128]-->[1,1,128]
        x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]

        x = self.to_latent(x)
        #再经过一层线性变换输出维度到num_classes
        #128->num_classes
        return self.mlp_head(x)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136