vit_train.py

import os
from PIL import Image
import torch
from torchvision import datasets, transforms
from utils.dataset import MNISTImageDataset
from utils.models import ViT
import torch.nn.functional as F
from tqdm import tqdm
from utils.backdoor import attack

backdoor = False
dataset_dir = "./dataset/MNIST"
wtype = 'hidden'
model_save_path = f"./models/vit_mnist_{"backdoor" if backdoor else "clean"}_{wtype}.pth"

batch_size = 128
lr = 3e-4
num_epochs = 10
img_width = 28
img_channels = 1
num_classes = 10
patch_size = 7
embedding_dim = 64
ff_dim = 2048
num_heads = 8
num_layers = 3
weight_decay = 1e-4


transform = transforms.Compose([
    transforms.ToTensor(),
])

train_backdoor_imgs = []
train_backdoor_labels = []
test_backdoor_imgs = []
test_backdoor_labels = []
if backdoor:
    backdoor_img_list, backdoor_label_list = attack(dataset_dir, transform, wtype=wtype)
    train_backdoor_imgs = backdoor_img_list[:500]
    train_backdoor_labels = backdoor_label_list[:500]
    test_backdoor_imgs = backdoor_img_list[500:]
    test_backdoor_labels = backdoor_label_list[500:]
    backdoor_test_dataset = MNISTImageDataset(root_dir=os.path.join(dataset_dir, "test"), 
                                              transform=transform, backdoor_img_list=test_backdoor_imgs, 
                                              backdoor_label_list=test_backdoor_labels, backdoor_only=True)
    backdoor_test_loader = torch.utils.data.DataLoader(backdoor_test_dataset, batch_size=batch_size, shuffle=False)

train_dataset = MNISTImageDataset(root_dir=os.path.join(
    dataset_dir, "train"), transform=transform, backdoor_img_list=train_backdoor_imgs, 
                                  backdoor_label_list=train_backdoor_labels)
test_dataset = MNISTImageDataset(root_dir=os.path.join(
    dataset_dir, "test"), transform=transform)


train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

device = "cuda:0" if torch.cuda.is_available() else "cpu"
print(f"{device=}")

model = ViT(
    img_width=img_width,
    img_channels=img_channels,
    patch_size=patch_size,
    d_model=embedding_dim,
    num_heads=num_heads,
    num_layers=num_layers,
    num_classes=num_classes,
    ff_dim=ff_dim,
).to(device)

print(model)



optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)

for epoch in range(num_epochs):
    losses = []
    total_train = 0
    correct_train = 0

    model.train()
    for img, label in tqdm(train_loader, desc=f"epoch-{epoch}"):
        img = img.to(device)
        label = label.to(device)

        pred = model(img)
        loss = F.cross_entropy(pred, label)

        pred_class = torch.argmax(pred, dim=1)
        correct_train += (pred_class == label).sum().item()
        total_train += pred.shape[0]

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        losses.append(loss.item())

    print(f"epoch-{epoch}: train loss:", sum(losses))
    print(f"epoch-{epoch}: train acc:", correct_train / total_train)

    model.eval()
    total = 0
    correct = 0
    with torch.no_grad():
        for img, label in test_loader:
            img = img.to(device)
            pred = torch.argmax(model(img), dim=1).cpu()

            correct += (pred == label).sum().item()
            total += pred.shape[0]

    print(f"epoch-{epoch}: test acc:", correct / total)
    
    backdoor_total = 0
    backdoor_correct = 0
    if backdoor:
        with torch.no_grad():
            for img, label in backdoor_test_loader:
                img = img.to(device)
                pred = torch.argmax(model(img), dim=1).cpu()

                backdoor_correct += (pred == label).sum().item()
                backdoor_total += pred.shape[0]
        
            print(f"epoch-{epoch}: backdoor test acc:", backdoor_correct / backdoor_total)
    

torch.save(model.state_dict(), model_save_path)