test.py

# Use machiraj_attks conda environment
import wandb
import sys
import argparse
import torch

# fix torch seed
torch.manual_seed(42)
# fix cuda seed
torch.cuda.manual_seed(42)
# fix cudnn seed for reproducibility
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

torch.cuda.empty_cache()
# sys.path.pop()
# wandb.login()
sys.path.insert(0, "..")
sys.path.append("../")
from misc.utils import *
from data.data import *
from models.model_loader import *
from eval.classic_eval import Evaluator
from eval.attack_eval import AttackEvaluator

# Defaults
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--threat_model",
        type=str,
        default="std",
        help="Threat model",
        choices=[
            "linf",
            "l2",
            "std",
            "untrained",
            "prime",
            "augmix",
            "xcit",
            "vit",
            "cc_sota",
            "da_augmix",
            "flc",
            "clad",
        ],
    )
    parser.add_argument(
        "--dataset",
        type=str,
        default="cifar10",
        help="Dataset",
        choices=[
            "cifar10",
            "imagenet",
            "cifar100",
            "imagenetc",
            "cifar10c",
            "cifar100c",
            "imagenet_9",
            "imagenet_9c",
            "custom_imagenet",
            "tiny_imagenet",
            "tiny_imagenetc",
        ],
    )

    parser.add_argument(
        "--number_of_iterations",
        type=int,
        default=100,
        help="Number of iterations for the Entropy attack",
    )


    parser.add_argument(
        "--pgd_iters", type=int, default=100, help="Number of iterations for PGD attack"
    )
    parser.add_argument(
        "--batch_size", default=256, type=int, help="Batch size for model evaluation",
    )

    parser.add_argument(
        "--lr", default=0.1, type=float, help="Learning rate for y_adv attack",
    )

    parser.add_argument(
        "--lambda_mse",
        default=0.0001,
        type=float,
        help="Regularization for y_adv attack",
    )

    parser.add_argument(
        "--save_mat",
        action="store_true",
        help="Save the y_quantize matrix for each image",
    )

    parser.add_argument(
        "--n_epochs",
        default=100,
        type=int,
        help="Number of iterations for y adversarial attack",
    )

    parser.add_argument(
        "--print_every",
        default=10,
        type=int,
        help="Print every n iterations for y adversarial attack",
    )

    parser.add_argument(
        "--retain_dc", action="store_true", help="Retain DC component for y_adv attack",
    )

    parser.add_argument(
        "--dc",
        default=0,
        type=int,
        help="Number of DC components to be kept to 1 in y_adv attack",
    )

    parser.add_argument(
        "--atk_type",
        choices=[
            "clean",
            "sharpen",
            "gaussian_blur",
            'box_blur',
            'non_local_means',
            'non_local_means_with_noise_estimation',
            # "ideal_edge_gray",
            # "entropy_l2",
            # "entropy_linf",
            "entropy_linf_single_scale",
            "entropy_anal",
            "pure",
            "time_anal",
            'nlm_time_anal',
            'sharpness_time_anal',
            "iteration_anal",
            'adaptive_eren',
            'adaptive_eren_median',
            'pca_train_anal',
            'adaptive_eren_estimate_p',
            'threshold_anal',
        ],
        default="pure",
        help="Correction type",
    )

    parser.add_argument(
        "--evaluate_on_corruptions",
        action="store_true",
        help="Evaluate on corruptions",
    )

    parser.add_argument(
        "--sharpness_factor",
        default=2.0,
        type=float,
        help="Sharpness factor for sharpening",
    )

    parser.add_argument(
        '--blur_factor',
        default=1.5,
        type=float,
        help="Sigma for Gaussian Blur",
    )

    parser.add_argument(
        "--threshold",
        default=0.,
        type=float,
        help="Threshold for entropy loss",
    )

    parser.add_argument(
        "--edge_kernel_size",
        default=3,
        type=int,
        help="Kernel size for edge detection",
    )

    parser.add_argument("--model_name", default="resnet50", type=str, help="Model name")
    parser.add_argument(
        "--parallel", action="store_true", help="Use DataParallel for model"
    )
    parser.add_argument("--verbose", action="store_true", help="Print losses in attack")
    parser.add_argument(
        "--lambda_prior",
        default=1,
        type=float,
        help="For attacks with prior the coefficient for the loss",
    )
    parser.add_argument(
        "--epsilon", default=8.0, type=float, help="For PGD Linf attack"
    )
    parser.add_argument(
        "--sim_loss",
        default="cosine",
        type=str,
        help="Type of loss for similarity of DCT",
    )
    parser.add_argument(
        "--kappa",
        default=0.9,
        type=float,
        help="Hinge Loss for Cosine Misclassification",
    )

    parser.add_argument(
        "--alpha",
        default=0.01,
        type=float,
        help="Step size for our entropy Linf/L2 attack",
    )

    parser.add_argument(
        '--test_anal',
        action="store_true",
        help="Run entropy analysis on test set",
    )

    args = parser.parse_args()
    param = vars(args)


    if param["atk_type"] == "entropy_anal" or param['atk_type'] == "time_anal":
        param["batch_size"] = 1

    config_wandb = dict(defense=param)
    param["adv_train"] = False
    run_name = run_name_generator(param)
    save_name = save_name_generator(param)

    param['shuffle'] = False
    param['train_mode'] = False

    logger = wandb.init(
        entity="harshitha-machiraju",
        project="EREN",
        reinit=True,
        name=run_name,
        config=config_wandb,
    )

    """
    Set Dataloaders and Model
    """

    if param["dataset"][-1] == "c":
        param["dataset"] = param["dataset"][:-1]



        # param["dataloader"] = None
    # elif param["atk_type"] == "entropy_anal" or param["dataset"][-1] != "c":
    data_loading = DataLoading(params=param)
    trainset, trainloader, testset, testloader = data_loading.get_data()
    if (param["atk_type"] == "entropy_anal" or  param["atk_type"] =='pca_train_anal') and param["dataset"]!="imagenet" and not param["test_anal"]:
        print("Train set loaded for entropy analysis")
        param["dataloader"] = trainloader
    else:
        print("Test set loaded")
        param["dataloader"] = testloader
        del trainset, trainloader, testset
    # else:
    #     raise NotImplementedError

    param["logger"] = logger

    model_loading = ModelLoader(params=param, device=device)
    net = model_loading.get_model()
    net = net.to(device)
    net = net.eval()

    print("Model loaded")

    """
	Evaluations and Attacks
	# """

    if param["atk_type"] == "clean":
        eval = Evaluator(param, device, net, logger)
        clean_acc = eval.clean_accuracy(testloader)

    else:
        eval = AttackEvaluator(device, net, param, logger)
        eval.attack_model()


    wandb.finish()


# CUDA_VISIBLE_DEVICES=0 python test.py --atk_type entropy --threat_model std --batch_size 32 --dataset cifar10 --model_name resnet50
# CUDA_VISIBLE_DEVICES=0 python test.py --atk_type sharpen --threat_model std --batch_size 32 --dataset cifar10 --model_name resnet50
# CUDA_VISIBLE_DEVICES=0 python test.py --atk_type sharpen --threat_model std --batch_size 64 --dataset imagenet --model_name resnet50
# CUDA_VISIBLE_DEVICES=3 python test.py --atk_type ideal_edge --threat_model std --batch_size 64 --dataset imagenet --model_name resnet50