File size: 2,921 Bytes
149f9ea
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import torch.nn as nn
import torch
import torch.optim as optim

class ViT(nn.Module):
    def __init__(self, image_size=28, patch_size=7, num_classes=10, dim=128, depth=6, heads=8, mlp_dim=256, dropout=0.1):
        super(ViT, self).__init__()

        assert image_size % patch_size == 0, 'image dimensions must be divisible by the patch size'
        num_patches = (image_size // patch_size) ** 2
        patch_dim = 1 * patch_size ** 2

        # 定义线性层将图像分块并映射到嵌入空间
        self.patch_embedding = nn.Linear(patch_dim, dim)

        # 位置编码
        # nn.Parameter是Pytorch中的一个类,用于将一个张量注册为模型的参数
        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches, dim))

        # Dropout层
        self.dropout = nn.Dropout(dropout)

        # Transformer编码器
        # 当 batch_first=True 时,输入和输出张量的形状为 (batch_size, seq_length, feature_dim)。当 batch_first=False 时,输入和输出张量的形状为 (seq_length, batch_size, feature_dim)。
        self.transformer = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(
                d_model=dim,
                nhead=heads,
                dim_feedforward=mlp_dim
                # batch_first=True
            ),
            num_layers=depth
        )

        # 分类头
        # nn.Identity()是一个空的层,它不执行任何操作,只是返回输入
        # self.to_cls_token = nn.Identity()
        # self.mlp_head = nn.Linear(dim, num_classes)
        self.mlp_head = nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, num_classes)
        )
    
    def forward(self, x):
        # x shape: [batch_size, 1, 28, 28]
        batch_size = x.size(0)
        x = x.view(batch_size, -1, 7*7)  # 将图像划分为7x7的Patch
        x = self.patch_embedding(x)  # [batch_size, num_patches, dim]
        x += self.pos_embedding  # 添加位置编码
        x = self.dropout(x)  # 应用Dropout

        x = x.permute(1, 0, 2)  # Transformer期望的输入形状:[seq_len, batch_size, embedding_dim]
        x = self.transformer(x)  # [序列长度, batch_size, dim]
        x = x.permute(1, 0, 2)  # 转回原来的形状:[batch_size, seq_len, dim]

        x = x.mean(dim=1)  # 对所有Patch取平均,x.mean(dim=1) 这一步是对所有 Patch 的特征向量取平均值,从而得到一个代表整个图像的全局特征向量。
        x = self.mlp_head(x)  # [batch_size, num_classes]
        return x

    # def forward(self, x):
    #     # x shape: (batch, 1, 28, 28)
    #     batch_size = x.shape[0]
    #     x = x.view(batch_size, -1, 7*7)
    #     x = self.patch_embedding(x)  # (batch, num_patches, dim)
    #     x = x + self.pos_embedding
    #     x = self.transformer(x)
    #     x = x.mean(dim=1)  # (batch, dim)
    #     x = self.mlp_head(x)
    #     return x