misc/vit.py

"""
ViT demo. 
The ViT model is from
    https://github.com/lucidrains/vit-pytorch/blob/main/vit_pytorch/vit.py
"""

import copy

import sys
import time

import matplotlib.pyplot as plt
import numpy as np 
import torch
import torchvision
import torchvision.transforms as transforms
from einops import rearrange, repeat
from einops.layers.torch import Rearrange
from torch import nn

sys.path.append("..")
import psgd

device = torch.device("cuda")


def set_seed(seed):
    # from chatgpt 
    np.random.seed(seed)                   # NumPy RNG
    torch.manual_seed(seed)                # PyTorch CPU RNG
    torch.cuda.manual_seed(seed)           # PyTorch GPU RNG (if using CUDA)

    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

set_seed(42)

"""
Prepare the dataset
"""
transform_train = transforms.Compose(
    [
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ]
)

transform_test = transforms.Compose(
    [
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ]
)

trainset = torchvision.datasets.CIFAR10(
    root="./data", train=True, download=True, transform=transform_train
)
train_loader = torch.utils.data.DataLoader(trainset, batch_size=100, shuffle=True)

testset = torchvision.datasets.CIFAR10(
    root="./data", train=False, download=True, transform=transform_test
)
test_loader = torch.utils.data.DataLoader(testset, batch_size=1000, shuffle=False)

"""
the ViT model is from
    https://github.com/lucidrains/vit-pytorch/blob/main/vit_pytorch/vit.py
"""
# helpers


def pair(t):
    return t if isinstance(t, tuple) else (t, t)


# classes


class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout=0.0):
        super().__init__()
        self.net = nn.Sequential(
            nn.LayerNorm(dim),
            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)


class Attention(nn.Module):
    def __init__(self, dim, heads=8, dim_head=64, dropout=0.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.norm = nn.LayerNorm(dim)

        self.attend = nn.Softmax(dim=-1)
        self.dropout = nn.Dropout(dropout)

        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):
        x = self.norm(x)

        qkv = self.to_qkv(x).chunk(3, dim=-1)
        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)
        attn = self.dropout(attn)

        out = torch.matmul(attn, v)
        out = rearrange(out, "b h n d -> b n (h d)")
        return self.to_out(out)


class Transformer(nn.Module):
    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.0):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.layers = nn.ModuleList([])
        for _ in range(depth):
            self.layers.append(
                nn.ModuleList(
                    [
                        Attention(dim, heads=heads, dim_head=dim_head, dropout=dropout),
                        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 self.norm(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.0,
        emb_dropout=0.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)"

        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.LayerNorm(patch_dim),
            nn.Linear(patch_dim, dim),
            nn.LayerNorm(dim),
        )

        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
        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.Linear(dim, num_classes)

    def forward(self, img):
        x = self.to_patch_embedding(img)
        b, n, _ = x.shape

        cls_tokens = repeat(self.cls_token, "1 1 d -> b 1 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)

        x = x.mean(dim=1) if self.pool == "mean" else x[:, 0]

        x = self.to_latent(x)
        return self.mlp_head(x)


# a tiny vit model
Net = ViT(
    image_size=32,
    patch_size=4,
    num_classes=10,
    dim=256,
    depth=4,
    heads=8,
    mlp_dim=256,
    dropout=0.1,
    emb_dropout=0.1,
)


"""
Test accuracy
"""


def test(net, data_loader):
    correct = 0
    total = 0
    with torch.no_grad():
        for _, (inputs, targets) in enumerate(data_loader):
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = net(inputs)
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()

    accuracy = correct / total
    return accuracy


"""
Now we compare Adam(W) (the default optimizer for transformer) and PSGD.
We align their settings, and the only difference is their preconditioners.  
"""
num_epochs = 100
plt.figure(figsize=(8, 4))
ax1 = plt.subplot(121)
ax2 = plt.subplot(122)
ax1.yaxis.tick_right()
ax2.yaxis.tick_right()

"""
Adam 
"""
net = copy.deepcopy(Net).to(device)
lr0 = 1e-3
opt = torch.optim.Adam(net.parameters(), lr=lr0)  # will aneal lr to lr0/num_epochs

TrainLoss = []
TestAcc = []
t0 = time.time()
for epoch in range(num_epochs):
    """train"""
    net.train()
    for _, (inputs, targets) in enumerate(train_loader):
        inputs, targets = inputs.to(device), targets.to(device)

        def closure():
            outputs = net(inputs)
            loss = nn.functional.cross_entropy(outputs, targets)
            return loss

        opt.zero_grad()
        loss = closure()
        loss.backward()
        opt.step()
        TrainLoss.append(loss.item())

    """test"""
    net.eval()
    test_acc = test(net, test_loader)
    TestAcc.append(test_acc)
    print(f"Adam, epoch {epoch + 1}, best test accuracy {max(TestAcc)}")

    opt.param_groups[0]["lr"] -= lr0 / num_epochs

total_time = time.time() - t0

SmoothedTrainLoss = np.convolve(TrainLoss, np.ones(100)/100)[99:-99]
ax1.plot(
    torch.arange(1, len(SmoothedTrainLoss) + 1).cpu() * total_time / len(SmoothedTrainLoss),
    SmoothedTrainLoss,
)
ax2.plot(
    torch.arange(1, len(TestAcc) + 1).cpu() * total_time / len(TestAcc),
    TestAcc,
)


"""
PSGD, Kron, gradient whitening, dQ={EQ, QE, QUAD, QEP}
"""
for dQ in ["Q0.5EQ1.5",]:
    net = copy.deepcopy(Net).to(device)
    
    lr0 = 1e-3  # keep the same as Adam
    opt = psgd.KronWhiten(
        net.parameters(),
        preconditioner_max_skew=2,
        lr_params=lr0,  # will aneal to lr0/num_epochs
        preconditioner_update_probability=1.0,  # anneal to 0.01
        momentum=0.9,
        dQ=dQ,
    )
    
    TrainLoss = []
    TestAcc = []
    
    t0 = time.time()
    for epoch in range(num_epochs):
        """train"""
        net.train()
        for _, (inputs, targets) in enumerate(train_loader):
            inputs, targets = inputs.to(device), targets.to(device)
    
            def closure():
                outputs = net(inputs)
                loss = nn.functional.cross_entropy(outputs, targets)
                return loss
    
            loss = opt.step(closure)
            TrainLoss.append(loss.item())
    
        """test"""
        net.eval()
        test_acc = test(net, test_loader)
        TestAcc.append(test_acc)
        print(f"PSGD, epoch {epoch + 1}, best test accuracy {max(TestAcc)}")
    
        opt.lr_params -= lr0 / num_epochs
        opt.preconditioner_update_probability = max(opt.preconditioner_update_probability * 0.1**0.1, 0.01)
    
    total_time = time.time() - t0
    
    SmoothedTrainLoss = np.convolve(TrainLoss, np.ones(100)/100)[99:-99]
    ax1.plot(
        torch.arange(1, len(SmoothedTrainLoss) + 1).cpu() * total_time / len(SmoothedTrainLoss),
        SmoothedTrainLoss,
    )
    ax2.plot(
        torch.arange(1, len(TestAcc) + 1).cpu() * total_time / len(TestAcc),
        TestAcc,
    )


ax1.set_xlabel("Wall time (s)", fontsize=6)
ax1.set_ylabel("Train loss", fontsize=6)
ax1.tick_params(labelsize=6)
ax1.legend(
    [
        "Adam",
        r"PSGD, $dQ=Q^{0.5}\mathcal{E}Q^{1.5}$",
    ],
    fontsize=7,
)
ax1.set_title("(a)", fontsize=7)

ax2.set_xlabel("Wall time (s)", fontsize=6)
ax2.set_ylabel("Test accuracy", fontsize=6)
ax2.tick_params(labelsize=6)
ax2.legend(
    [
        "Adam",
        r"PSGD, $dQ=Q^{0.5}\mathcal{E}Q^{1.5}$",
    ],
    fontsize=7,
)
ax2.set_title("(b)", fontsize=7)

plt.savefig("vit_adam_vs_psgd.svg")
plt.savefig("vit_adam_vs_psgd.eps")
plt.show()