add vit

Author: yixu.cui

models/vit/vit_mnist.py

@@ -0,0 +1,311 @@
+import math
+import os
+import time
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torch.optim.lr_scheduler as lr_scheduler
+import torchvision
+from einops import rearrange
+from torch import nn
+from torch import optim
+from torchvision import datasets
+from torchvision import transforms
+from torchvision.transforms import ToTensor
+
+
+plt.ion()  # interactive mode
+
+torch.manual_seed(42)
+DOWNLOAD_PATH = "/share-global/yixu.cui/datas/mnist"
+DOWNLOAD_DATA_PATH = "/share-global/yixu.cui/datas/"
+BATCH_SIZE_TRAIN = 256 * 8 * 2
+BATCH_SIZE_TEST = 256 * 8 * 2 * 2
+
+
+# device: GPU
+os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1, 2, 3, 4, 5, 6, 7"
+os.environ["CUDA_VISIBLE_DEVICES"] = "4, 5, 6, 7"
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = (
+    "cuda"
+    if torch.cuda.is_available()
+    else "mps"
+    if torch.backends.mps.is_available()
+    else "cpu"
+)
+
+gpu_num = torch.cuda.device_count()
+device_ids = [i for i in range(gpu_num)]
+
+
+# 残差模块，放在每个前馈网络和注意力之后
+class Residual(nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, x, **kwargs):
+        return self.fn(x, **kwargs) + x
+
+
+# layernorm归一化,放在多头注意力层和激活函数层。用绝对位置编码的BERT，layernorm用来自身通道归一化
+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):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, dim)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+# 多头注意力层，多个自注意力连起来。使用qkv计算
+class Attention(nn.Module):
+    def __init__(self, dim, heads=8):
+        super().__init__()
+        self.heads = heads
+        self.scale = dim**-0.5
+        self.to_qkv = nn.Linear(dim, dim * 3, bias=False)
+        self.to_out = nn.Linear(dim, dim)
+
+    def forward(self, x, mask=None):
+        b, n, _, h = *x.shape, self.heads
+        qkv = self.to_qkv(x)
+        q, k, v = rearrange(qkv, "b n (qkv h d) -> qkv b h n d", qkv=3, h=h)
+        dots = torch.einsum("bhid,bhjd->bhij", q, k) * self.scale
+        if mask is not None:
+            mask = F.pad(mask.flatten(1), (1, 0), value=True)
+            assert mask.shape[-1] == dots.shape[-1], "mask has incorrect dimensions"
+            mask = mask[:, None, :] * mask[:, :, None]
+            dots.masked_fill_(~mask, float("-inf"))
+            del mask
+        attn = dots.softmax(dim=-1)
+        out = torch.einsum("bhij,bhjd->bhid", attn, v)
+        out = rearrange(out, "b h n d -> b n (h d)")
+        out = self.to_out(out)
+        return out
+
+
+class Transformer(nn.Module):
+    def __init__(self, dim, depth, heads, mlp_dim):
+        super().__init__()
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(
+                nn.ModuleList(
+                    [
+                        Residual(PreNorm(dim, Attention(dim, heads=heads))),
+                        Residual(PreNorm(dim, FeedForward(dim, mlp_dim))),
+                    ]
+                )
+            )
+
+    def forward(self, x, mask=None):
+        for attn, ff in self.layers:
+            # print(f"batch size: {x.shape[0]}")  # debug to locate how many img in per GPU
+            x = attn(x, mask=mask)
+            x = ff(x)
+        return x
+
+
+# 将图像切割成一个个图像块,组成序列化的数据输入Transformer执行图像分类任务。
+class ViT(nn.Module):
+    def __init__(
+        self,
+        *,
+        image_size,
+        patch_size,
+        num_classes,
+        dim,
+        depth,
+        heads,
+        mlp_dim,
+        channels=3,
+        mlp_drop_ratio=0.0,
+    ):
+        super().__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 = channels * patch_size**2
+        self.patch_size = patch_size
+        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
+        self.patch_to_embedding = nn.Linear(patch_dim, dim)
+        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
+        self.transformer = Transformer(dim, depth, heads, mlp_dim)
+        self.to_cls_token = nn.Identity()
+        self.mlp_head = nn.Sequential(
+            nn.Linear(dim, mlp_dim),
+            nn.Dropout(p=mlp_drop_ratio),
+            nn.GELU(),
+            nn.Linear(mlp_dim, num_classes),
+            # nn.Dropout(p=mlp_drop_ratio),
+        )
+
+    def forward(self, img, mask=None):
+        p = self.patch_size
+        x = rearrange(img, "b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1=p, p2=p)
+        x = self.patch_to_embedding(x)
+        cls_tokens = self.cls_token.expand(img.shape[0], -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+        x += self.pos_embedding
+        x = self.transformer(x, mask)
+        x = self.to_cls_token(x[:, 0])
+        return self.mlp_head(x)
+
+
+###
+def train_epoch(model, optimizer, data_loader, loss_history):
+    total_samples = len(data_loader.dataset)
+    model.train()
+
+    for i, (data, target) in enumerate(data_loader):
+        optimizer.zero_grad()
+        output = F.log_softmax(model(data.to(device)), dim=1)
+        loss = F.nll_loss(output, target.to(device))
+        loss.backward()
+        optimizer.step()
+
+        if i % 128 == 0:
+            print(
+                "["
+                + "{:5}".format(i * len(data))
+                + "/"
+                + "{:5}".format(total_samples)
+                + " ("
+                + "{:3.0f}".format(100 * i / len(data_loader))
+                + "%)]  Loss: "
+                + "{:6.4f}".format(loss.item())
+            )
+            loss_history.append(loss.item())
+
+
+def evaluate(model, data_loader, loss_history):
+    model.eval()
+
+    total_samples = len(data_loader.dataset)
+    correct_samples = 0
+    total_loss = 0
+
+    with torch.no_grad():
+        for data, target in data_loader:
+            output = F.log_softmax(model(data.to(device)), dim=1)
+            loss = F.nll_loss(output, target.to(device), reduction="sum")
+            _, pred = torch.max(output, dim=1)
+
+            total_loss += loss.item()
+            correct_samples += pred.eq(target.to(device)).sum()
+
+    avg_loss = total_loss / total_samples
+    loss_history.append(avg_loss)
+    print(
+        "\nAverage test loss: "
+        + "{:.4f}".format(avg_loss)
+        + "  Accuracy:"
+        + "{:5}".format(correct_samples)
+        + "/"
+        + "{:5}".format(total_samples)
+        + " ("
+        + "{:14.2f}".format(100.0 * correct_samples / total_samples)
+        + "%)\n"
+    )
+
+
+if __name__ == "__main__":
+    # mnist
+    transform_mnist = torchvision.transforms.Compose(
+        [
+            torchvision.transforms.ToTensor(),
+            torchvision.transforms.Normalize((0.1307,), (0.3081,)),
+        ]
+    )
+
+    train_set = torchvision.datasets.MNIST(
+        DOWNLOAD_PATH, train=True, download=True, transform=transform_mnist
+    )
+    train_loader = torch.utils.data.DataLoader(
+        train_set, batch_size=BATCH_SIZE_TRAIN, shuffle=True
+    )
+
+    test_set = torchvision.datasets.MNIST(
+        DOWNLOAD_PATH, train=False, download=True, transform=transform_mnist
+    )
+    test_loader = torch.utils.data.DataLoader(
+        test_set, batch_size=BATCH_SIZE_TEST, shuffle=True
+    )
+
+    for X, y in test_loader:
+        print(f"Shape of X [N, C, H, W]: {X.shape}")
+        print(f"Shape of y: {y.shape} {y.dtype}")
+        break
+
+    EPOCHS_NUM = 100
+    # start_time = time.time()
+
+    """
+    patch大小为 7x7（对于 28x28 图像，这意味着每个图像 4 x 4 = 16 个patch）、10 个可能的目标类别（0 到 9）和 1 个颜色通道（因为图像是灰度）。
+    在网络参数方面，使用了 64 个单元的维度，6 个 Transformer 块的深度，8 个 Transformer 头，MLP 使用 128 维度。
+    """
+    model = ViT(
+        image_size=28,
+        patch_size=7,
+        num_classes=10,
+        channels=1,
+        dim=64,
+        depth=6,
+        heads=8,
+        mlp_dim=128,
+        mlp_drop_ratio=0.5,
+    )
+    # print(f"No device:\n{model}")
+    # 模型并行化
+    if torch.cuda.device_count() > 1:
+        print("Let's use", torch.cuda.device_count(), "GPUs!")
+        model = nn.DataParallel(model)
+        # model = nn.parallel.DistributedDataParallel(model)
+
+    model = model.to(device)
+    # print(f"On device:\n{model}")
+
+    loss_fn = (
+        nn.CrossEntropyLoss()
+    )  # Ref: https://pytorch.org/tutorials/beginner/basics/quickstart_tutorial.html
+
+    optimizer = optim.Adam(model.parameters(), lr=0.008)
+    # Scheduler https://arxiv.org/pdf/1812.01187.pdf
+    # lf = (
+    #     lambda x: ((1 + math.cos(x * math.pi / args.epochs)) / 2) * (1 - args.lrf)
+    #     + args.lrf
+    # )  # cosine
+    # scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
+
+    train_loss_history, test_loss_history = [], []
+    for epoch in range(1, EPOCHS_NUM + 1):
+        print("Epoch:", epoch)
+        start_time = time.time()
+        train_epoch(model, optimizer, train_loader, train_loss_history)
+        evaluate(model, test_loader, test_loss_history)
+        print(
+            "This EPOCH takes time:",
+            "{:5.2f}".format(time.time() - start_time),
+            "seconds",
+        )
+
+    # print("Execution time:", "{:5.2f}".format(time.time() - start_time), "seconds")