t2t_train.py

import random
from tqdm import tqdm
from utils.dataset import CIFAR10ImageDataset
from torchvision import transforms
import toml
import torch
import os
from utils.backdoor import attack_cf
from utils.models import ImprovedT2TViT
import torch.nn.functional as F

backdoor = True
dataset_dir = "./dataset/CIFAR10"
wtype = 'content'
model_save_path = f"./models/t2t_cifar10_{"backdoor" if backdoor else "clean"}_{wtype}.pth"




config = toml.load(f"./config/cifar10.toml")


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

train_transform = transforms.Compose([
    transforms.RandomCrop(32, padding=4),  
    transforms.RandomHorizontalFlip(),      
    transforms.ToTensor(),                  
    transforms.Normalize(
        mean=[0.4914, 0.4822, 0.4465],
        std=[0.2023, 0.1994, 0.2010]
    )
])
train_backdoor_imgs = []
train_backdoor_labels = []
test_backdoor_imgs = []
test_backdoor_labels = []
if backdoor:
    backdoor_img_list, backdoor_label_list = attack_cf(dataset_dir, transform, wtype=wtype, dataset="cifar10")
    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 = CIFAR10ImageDataset(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=config["train"]["batch_size"], shuffle=False)
    

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

print(f"train size:{len(train_dataset)}, test size:{len(test_dataset)}, backdoor test size:{len(backdoor_test_dataset)}")

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=config["train"]["batch_size"], shuffle=True)

test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=config["train"]["batch_size"], shuffle=False)


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

model = ImprovedT2TViT(
    img_size=config["model"]["img_size"],
    in_channels=config["model"]["in_channels"],
    patch_size=config["model"]["patch_size"],
    d_model=config["model"]["d_model"],
    num_heads=config["model"]["num_heads"],
    num_layers=config["model"]["num_layers"],
    num_classes=config["model"]["num_classes"],
    ff_dim=config["model"]["ff_dim"],
).to(device)

print(model)



optimizer = torch.optim.AdamW(model.parameters(), lr=config["train"]["lr"], weight_decay=config["train"]["weight_decay"])

for epoch in range(config["train"]["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
    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)