initial commit from desktop

README.md

@@ -0,0 +1 @@
+This is a MNIST classifier based on vision transformer.

vit.py

@@ -0,0 +1,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

vit_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:223c6c32c2a9d4c274b09c35ef089b358ee7cf1729b9d939fca898db5765dcdb
+size 3248655