1.一脚踹进ViT——Pytorch搭建ViT框架_pytorch vit

一脚踹进ViT——Pytorch搭建ViT框架

本系列根据百度飞浆Paddle教程，学习整理后的博客，本文主要使用pytorch对残差网络ResNet18进行实现，首先对代码以及结构搭建进行熟悉，进而介绍简单的机器学习以及tensor使用，最后实现ViT的基本框架，请各位仔细食用！

1.ResNet18的实现

第一部分为了对经典残差网络进行复现，后续ViT中也将残差思想多次运用

1.1 首先，搭建最简单的网络框架

mport torch
import torch.nn as n

def main():
    # 定义一个tensor，四个数分别是 batch_size=4，channel=3, 图像宽、高都是32，32
    t = torch.randn([4,3,32,32])
    # 定义一个ResNet模型类
    model = ResNet()
    # 打印模型参数供查看
    print(model)
    # 将输入传入模型得到结果
    out = model(t)
    print(out.shape)

if __name__=="__main__":
    main()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

1.2 创建ResNet类

定义好框架后，需要定义ResNet类了，需要定义 init函数和 forward前向传播函数，查看图中结构，输入后经过7×7的卷积，我们实现用3×3的卷积核代替，结构用到的无非就是卷积，最大池化，多个卷积块（其中卷积、BN、激活）、平均池化，因为卷积中仅输出通道不同，以及有多个block堆叠，我们可以定义一个生成卷积层的函数 以及 创建Block的函数 。

对于我们网络的输入通道数设置：in_dim=64，最终假设同MNIST数据集num_classes=10，之前传入的tensor尺寸是28×28的，我们模拟该过程。

class ResNet(nn.Module):
    def __init__(self,in_dim=64,num_classes=10):                    
        # 继承父类，记得写
        super().__init__()
        # 定义所需内容
    def forward(self,x):
        # 对输入进行处理
        return x
    def _make_layer(self,dim,n_blocks,stride):
        # 构建卷积模块
        return nn.Sequential(*layer_list)
1
2
3
4
5
6
7
8
9
10
11

接着我们对于每一步进行补充即可

1.2.1 初始化函数定义所需操作

    def __init__(self,in_dim=64,num_classes=10):                    
        super().__init__()
        self.in_dim = in_dim
        self.conv1 = nn.Conv2d(in_channels=3,
                               out_channels=in_dim,
                               kernel_size=3,
                               stride=1,
                               padding=1,
                               bias=False)
        self.bn1 = nn.BatchNorm2d(in_dim)
        self.relu = nn.ReLU()
        # 定义blocks，第一个卷积层，输入输出维度不变，后三层其中第一块将图像减半，第二次卷积图像大小不变，每层输出通道数加倍
        self.layer1 = self._make_layer(dim=64, n_blocks=2,stride=1)
        self.layer2 = self._make_layer(dim=128, n_blocks=2,stride=2)
        self.layer3 = self._make_layer(dim=256, n_blocks=2,stride=2)
        self.layer4 = self._make_layer(dim=512, n_blocks=2,stride=2)

        # head layer
        self.avgpool = nn.AdaptiveAvgPool2d(1)
        self.classifier = nn.Linear(512,num_classes)  #最终分类


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

1.2.2 forward构建模块

    def forward(self,x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.avgpool(x)
        # 对于后面维度进行拉直
        x = x.flatten(1)
        x = self.classifier(x)
1
2
3
4
5
6
7
8
9
10
11
12

1.3.3 _make_layer函数

    def _make_layer(self,dim,n_blocks,stride):
        
        # 定义一个list，存放构建的block，dim是out dim
        layer_list = []
        layer_list.append(Block(self.in_dim,dim,stride))
        self.in_dim = dim
        for i in range(1, n_blocks):
            layer_list.append(Block(self.in_dim, dim, stride=1))
        # 使用* 将list中所有元素取出来 构建层
        return nn.Sequential(*layer_list)
1
2
3
4
5
6
7
8
9
10

1.3 Block块

每个Block块由两次conv bn 以及一个激活函数组成， 对于stride=2 或者 输入维度和输出维度不一样的时候，说明是 第一次构建块或者残差实现时，则加一个下采样层，得到输出维度，否则进行一个恒等映射Indentity()，这个类其实很简单，将输入作为结果进行输出。

class Block(nn.Module):
    def __init__(self,in_dim,out_dim,stride):
        super().__init__()
        self.conv1 = nn.Conv2d(in_dim,out_dim,3 ,stride=stride,padding=1,bias=False)
        self.bn1 = nn.BatchNorm2d(out_dim)
        self.conv2= nn.Conv2d(out_dim,out_dim,3 ,stride=1,padding=1,bias=False)
        self.bn2 = nn.BatchNorm2d(out_dim)
        self.relu = nn.ReLU()

        if stride == 2 or in_dim != out_dim:
            self.downsample = nn.Sequential(*[
                nn.Conv2d(in_dim,out_dim,1,stride=stride),
                nn.BatchNorm2d(out_dim)])
        else:
            self.downsample = Indentity()

    def forward(self,x):
        # 残差实现，将x保留作为h，最终将维度统一后与经过特征提取后的结果相加
        h = x
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.bn2(x)
        identity = self.downsample(h)
        x = x+identity
        x = self.relu(x)
        return 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

这样就构建完成了，最终将Indentity()类一补充就完成了，是不是很清晰呢？

class Indentity(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self,x):
        return x
1
2
3
4
5
6

看看结果吧？

网络结构输出可以自行输出看一下，我这里放一部分，因为我们的网络太大了！哈哈

输出的Shape：torch.Size([4, 10])

是不是就得到了4个batch，每个中十类的值，最终接个softmax就可以分类了。

2.机器学习

正常情况在数学情况中，2x+1=y ，求x和y；

而机器学习是，ax+b=y，给定（x,y），让你求最优的a和b，而线性函数表示范围比较小，用二次函数能拟合更多情况，那是不是越复杂的函数他就可以表示越多种情况呢?

这个函数就能叫做model，model可以是任意复杂的函数，我们机器学习就是来找最优的参数拟合函数

3.tensor

简单介绍一下pytorch中 tensor的简单构建和操作

from PIL import Image
import numpy as np
import torch

torch.device('cpu')

def main():

    # 1. 创建一个tensor
    t = torch.zeros([3,3])
    print(t)

    # 2. 创建一个随机tensor
    t = torch.randn([5,3])
    print(t)

    # 3.从图片中创建一个tensor
        # 随机加载一个 28×28的图片，遍历输出
    img = np.array(Image.open('ok.jpg'))
    for i in range(28):
        for j in range(28):
           print(f'{img[i,j]:03} ',end='')
        print()
    t = torch.tensor(img,dtype = torch.float32)

    # 4. 打印tensor类型，和tensor的 dtype（里面值的类型）
    print(type(t))
    print(t.dtype)

    # print(t.transpose(1,0))

    # 5.chunk 切分
    t = torch.randint(0,10,[5,15])
    print(t)
    qkv = t.chunk(3,-1)
    q,k,v = qkv
    print(q)

if __name__ == '__main__':
    main()
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

根据代码中五个部分分别查看不同输出来看函数的使用

第一个输出为：

tensor([[0., 0., 0.], [0., 0., 0.], [0., 0., 0.]])

第二个输出为：

tensor([[ 0.8475, -0.6581, -0.2290], [ 0.4138, 0.2462, -1.1374], [-0.1824, -1.6705, -0.4217], [ 0.1481, -0.7515, -1.7304], [ 0.2509, -0.1405, 0.2673]])

第三个输出 太大了，如果读取的是MNIST数据集的话，可以根据像素值，远看表示的哪个数字，大家可以试试

第四个输出：

<class 'torch.Tensor'> torch.float32

type来表示t的类型是 Tensor，而 dtype是其中的值的类型

第五个输出：

可以看到原始的t被切成三份q,k,v，竖着切了三分，由图可见

4.ViT中 Patch Embedding

ViT中最重要的就是将图像切为Patch块，排成序列后，通过线性映射，通过image tokens输入网络最终得到特征

Patch Embedding就作为 image tokens输入网络，进行多个encoder进行特征提取，每一个image token/patch embedding 输入到 transformer中它还会单独输出一个tensor
若输入是 7*7=49个tensor，输出还是49个tensor，那怎么做分类呢？
可以用avgpool取平均，放入head中

接下来就实现Patch Embedding 同样是类似刚ResNet构建思路，搭建框架，主函数中输入一个28×28的图像，然后Patch Embedding就会将图像切为Patch，得到Image Patches后，进行线性投射（可以用卷积来实现，卷积核大小与Patch块大小一样，步长为卷积块大小），可以得到Patch Embeddings，我们embed_dim设置为1，最终得到的就是单个序列

import torch
import torch.nn as nn
from PIL import Image
import numpy as np

class PatchEmbedding(nn.Module):
    def __init__(self,image_size, patch_size, in_channels, embed_dim ,dropout=0.):
        super().__init__()
        self.patch_embed = nn.Conv2d(in_channels,
                                     embed_dim,
                                     kernel_size=patch_size,
                                     stride=patch_size,
                                     bias=False)
        self.patch_embed.weight = nn.Parameter(nn.init.constant_(self.patch_embed.weight,1.0))
        self.dropout = nn.Dropout(p=dropout)



    def forward(self,x):
        # x: [1, 1, 28, 28]
        x = self.patch_embed(x)
        # x: [n, embed_dim, h', w']
        x = x.flatten(2)   #后两个维度拉直[n, embed_dim, h'*w']
        x = x.permute(0, 2, 1)  #[n,  h'*w', embed_dim]
        x = self.dropout(x)
        return x

def main():
    # 1. 加载图像并且转化为tensor
    img = Image.open('ok.jpg')
    img = np.array(img)
    for i in range(28):
        for j in range(28):
            print(f'{img[i, j]:03} ', end='')
        print()

    sample = torch.tensor(img,dtype=torch.float32)
    # 模拟一个batch的数据
    sample = sample.reshape([1,1,28,28])
    print(sample.shape)    #[1,1,28,28]

    # 2. Patch Embedding
    patch_embed = PatchEmbedding(image_size =28, patch_size=7, in_channels=1,embed_dim=1)
    out = patch_embed(sample)
    print(out.shape)
    print(out)
    for i in range(0,28, 7):
        for j in range(0, 28, 7):
            # 固定卷积核中weight权重为1，卷积后输出与直接将对应位置数字相加后得到的相同
            print(torch.sum(sample[0,0,i:i+7, j:j+7]).numpy().item())

if __name__ == '__main__':
    main()
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

可以看到经过Patch Embedding之后得到的值，与（进行weight全为1）卷积的输出匹配，所以将图像切成了16个 7×7大小的块，tensor的shape为：torch.Size([1, 16, 1]) ，分别是batch_size大小，patch的长度，以及embed_dim维度大小，我们设置的为1。

接下来再加个MLP层，初始化 定义参数embed_dim，以及mlp_ratio进行输出维度大小的扩大，使用GELU激活函数

class Indentity(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self,x):
        return x

class Mlp(nn.Module):
    def __init__(self,embed_dim, mlp_ratio=4.0, dropout=0.):
        super().__init__()
        self.fc1  = nn.Linear(embed_dim, int(embed_dim* mlp_ratio))
        self.fc2  = nn.Linear(int(embed_dim* mlp_ratio),embed_dim)
        self.act  = nn.GELU()
        self.dropout = nn.Dropout(p=dropout)


    def forward(self,x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.dropout(x)
        return 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

主函数中加入MLP的调用

    # 3.MLP
    mlp = Mlp(1)
    out = mlp(out)
    print('out shape=',out.shape)
1
2
3
4

输入维度 torch.Size([1, 1, 28, 28])

patch embedding后 torch.Size([1, 16, 1])

经过MLP输出维度out shape= torch.Size([1, 16, 1])

除了MSA没有实现，其实ViT大致已经可以实现了，我们三步走：

1. 用卷积操作做 patch Embedding 得到了 Image Tokens

2. 实现ViT class，其中实现了五层的Encoder

3. 最后用head 头做分类， 一般是Liner接一个Softmax层，本次实现中省掉了softmax，直接拿linear做分类

有了这些在forward中：图像做patch embedding得到了Image Tokens，分别五次前向Encoder，之后经过Linear层

而 Encoder部分，MSA用 TODO代替，暂时先不做，用residual块进行连接，进行LayerNorm，实现了MLP（FFN），之后又实现了residual block，再实现一次LayerNorm。将它复制N次，构成多个Encoder。

定义一个ViT类

import torch
import torch.nn as nn

class Indentity(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self,x):
        return x

class Mlp(nn.Module):
    def __init__(self,embed_dim, mlp_ratio=4.0, dropout=0.):
        super().__init__()
        self.fc1  = nn.Linear(embed_dim, int(embed_dim* mlp_ratio))
        self.fc2  = nn.Linear(int(embed_dim* mlp_ratio),embed_dim)
        self.act  = nn.GELU()
        self.dropout = nn.Dropout(p=dropout)

    def forward(self,x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.dropout(x)
        return x

class PatchEmbedding(nn.Module):
    def __init__(self,image_size, patch_size, in_channels, embed_dim ,dropout=0.):
        super().__init__()
        self.patch_embed = nn.Conv2d(in_channels,
                                     embed_dim,
                                     kernel_size=patch_size,
                                     stride=patch_size,
                                     bias=False)
        self.patch_embed.weight = nn.Parameter(nn.init.constant_(self.patch_embed.weight,1.0))
        self.dropout = nn.Dropout(p=dropout)

    def forward(self,x):
        # x: [1, 1, 28, 28]
        x = self.patch_embed(x)
        # x: [n, embed_dim, h', w']
        x = x.flatten(2)   #[n, embed_dim, h'*w']
        x = x.permute(0, 2, 1)  #[n,  h'*w', embed_dim]
        x = self.dropout(x)
        return x

class Encoder(nn.Module):
    def __init__(self,embed_dim):
        super().__init__()
        self.attn = Indentity()  #TODO
        self.attn_norm = nn.LayerNorm(embed_dim)
        self.mlp = Mlp(embed_dim)
        self.mlp_norm = nn.LayerNorm(embed_dim)

    def forward(self,x):
        h = x
        x = self.attn_norm(x)
        x = self.attn(x)
        x = x+h

        h = x
        x = self.mlp_norm(x)
        x = self.mlp(x)
        x = x+h

        return x

class ViT(nn.Module):
    def __init__(self):
        super().__init__()
        self.patch_embed = PatchEmbedding(224, 7, 3, 16)
        # 存五个Encoder
        layer_list = [Encoder(16) for i in range(5)]
        self.encoders = nn.ModuleList(layer_list)
        self.head = nn.Linear(16, 10)  #10:num_classes
        self.avgpool = nn.AdaptiveAvgPool1d(1)

    def forward(self, x):
        x = self.patch_embed(x)
        for encoder in self.encoders:
            x = encoder(x)

        # layernorm
        # [n, h'*w', c]
        x = x.permute(0, 2, 1)
        x = self.avgpool(x)    # [n , c , 1]
        x = x.flatten(1)  # [n, c]
        x = self.head(x)
        return x

def main():
    t = torch.randn([4, 3, 224, 224])
    model = ViT()
    out = model(t)

    print('out shape=',out.shape)

if __name__ == '__main__':
    main()
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

结果：out shape= torch.Size([4, 10])

基本ViT的实现框架形成了，但MSA（Mult-Head Self-Attention）自注意力机制的实现下节进行实现