import torch
import torch.nn as nn
import sys
sys.path.append('../')
import TransformerConfigs_pretrain as configs
from TransformerModels_pretrain import ViTModel_custom, ViTForImageClassification

sys.path.append('../../')
from load_data import load_tiny, GetCIFAR100Validation, GetCIFAR10Validation
from Evaluations import test_vit

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

Using device:  cuda (NVIDIA A100-PCIE-40GB)

# load dataset and model weights
dataset = 'TinyImagenet'
num_labels = 200
model_arch = 'ViT-16'
config = configs.get_b16_config()
test_loader = load_tiny()
model = ViTModel_custom(config=config)
model = ViTForImageClassification(config, model, num_labels)
filename = "../results/{}/{}/weights.pth".format(model_arch, dataset)
model.load_state_dict(torch.load(filename), strict=False)
model = nn.DataParallel(model).cuda()
model.eval()

DataParallel(
  (module): ViTForImageClassification(
    (vit): ViTModel_custom(
      (embeddings): ViTEmbeddings(
        (patch_embeddings): ViTPatchEmbeddings(
          (projection): Conv2d(3, 768, kernel_size=(16, 16), stride=(16, 16))
        )
        (dropout): Dropout(p=0.0, inplace=False)
      )
      (encoder): ViTEncoder(
        (layer): ModuleList(
          (0-11): 12 x ViTLayer(
            (attention): ViTAttention(
              (attention): ViTSelfAttention(
                (query): Linear(in_features=768, out_features=768, bias=True)
                (key): Linear(in_features=768, out_features=768, bias=True)
                (value): Linear(in_features=768, out_features=768, bias=True)
                (dropout): Dropout(p=0.0, inplace=False)
              )
              (output): ViTSelfOutput(
                (dense): Linear(in_features=768, out_features=768, bias=True)
                (dropout): Dropout(p=0.0, inplace=False)
              )
            )
            (intermediate): ViTIntermediate(
              (dense): Linear(in_features=768, out_features=3072, bias=True)
            )
            (output): ViTOutput(
              (dense): Linear(in_features=3072, out_features=768, bias=True)
              (dropout): Dropout(p=0.0, inplace=False)
            )
            (layernorm_before): LayerNorm((768,), eps=1e-12, elementwise_affine=True)
            (layernorm_after): LayerNorm((768,), eps=1e-12, elementwise_affine=True)
          )
        )
      )
      (layernorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)
      (pooler): ViTPooler(
        (dense): Linear(in_features=768, out_features=768, bias=True)
        (activation): Tanh()
      )
    )
    (dropout): Dropout(p=0.1, inplace=False)
    (classifier): Linear(in_features=768, out_features=200, bias=True)
  )
)

No attack

_, adv_acc = test_vit(model=model, test_loader=test_loader, device=device)

Test loss: 0.64
Test accuracy: 85.03%

PGD

adv_filepath = "../results/{}/{}/adv_results/".format(model_arch, dataset)
advLoader = torch.load(adv_filepath+'PGD_advLoader.pth')
_, adv_acc = test_vit(model=model, test_loader=advLoader, device=device)

Test loss: 11.27
Test accuracy: 0.05%

CW

advLoader = torch.load(adv_filepath+'CW_advLoader.pth')
_, adv_acc = test_vit(model=model, test_loader=advLoader, device=device)

Test loss: 12.28
Test accuracy: 0.10%

SGM

advLoader = torch.load(adv_filepath+'SGM_advLoader.pth')
_, adv_acc = test_vit(model=model, test_loader=advLoader, device=device)

Test loss: 3.19
Test accuracy: 34.74%

PatchFool

advLoader = torch.load(adv_filepath+'PatchFool_advLoader.pth')
_, adv_acc = test_vit(model=model, test_loader=advLoader, device=device)

Test loss: 6.12
Test accuracy: 14.55%