eval_trades.py

"""
CIFAR-10 evaluation script from TRADES.

Reference:
    https://github.com/yaodongyu/TRADES/blob/master/evaluate_attack_cifar10.py
"""
from __future__ import print_function
import os
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from torch.autograd import Variable
import torch.optim as optim
from torchvision import datasets, transforms
from wideresnet import *
import numpy as np


parser = argparse.ArgumentParser(description='PyTorch CIFAR Attack Evaluation')
parser.add_argument('--no-cuda', action='store_true', default=False,
                    help='disables CUDA training')
parser.add_argument('--epsilon', default=0.031,
                    help='perturbation')
parser.add_argument('--model-path',
                    default='./models/model_cifar_wrn.pt',
                    help='model for white-box attack evaluation')
parser.add_argument('--data-attak-path',
                    default='./attacks/cifar10_X_adv.npy',
                    help='adversarial data for white-box attack evaluation')
parser.add_argument('--data-path',
                    default='./attacks/cifar10_X.npy',
                    help='data for white-box attack evaluation')
parser.add_argument('--target-path',
                    default='./attacks/cifar10_Y.npy',
                    help='target for white-box attack evaluation')

args = parser.parse_args()

# settings
use_cuda = not args.no_cuda and torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}


def image_check(min_delta, max_delta, min_image_adv, max_image_adv):
    valid = 1.0
    if min_delta < - args.epsilon:
        print(f'{min_delta} < -{args.epsilon}')
        valid -= 2.0
    elif max_delta > args.epsilon:
        print(f'{max_delta} > {args.epsilon}')
        valid -= 2.0
    elif min_image_adv < 0.0:
        print(f'{min_image_adv} < 0')
        valid -= 2.0
    elif max_image_adv > 1.0:
        print(f'{max_image_adv} > 1')
        valid -= 2.0

    if valid > 0.0:
        return True
    else:
        return False


def eval_adv_test_whitebox(model, device, X_adv_data, X_data, Y_data):
    """
    evaluate model by white-box attack
    """
    model.eval()
    robust_err_total = 0

    with torch.no_grad():
        for idx in range(len(Y_data)):
            # load original image
            image = np.array(np.expand_dims(X_data[idx], axis=0), dtype=np.float32)
            image = np.transpose(image, (0, 3, 1, 2))
            # load adversarial image
            image_adv = np.array(np.expand_dims(X_adv_data[idx], axis=0), dtype=np.float32)
            image_adv = np.transpose(image_adv, (0, 3, 1, 2))
            # load label
            label = np.array(Y_data[idx], dtype=np.int64)

            # check bound
            image_delta = image_adv - image
            min_delta, max_delta = image_delta.min(), image_delta.max()
            min_image_adv, max_image_adv = image_adv.min(), image_adv.max()
            valid = image_check(min_delta, max_delta, min_image_adv, max_image_adv)
            if not valid:
                print('not valid adversarial image')
                break

            # transform to torch.tensor
            data_adv = torch.from_numpy(image_adv).to(device)
            target = torch.from_numpy(label).to(device)

            # evluation
            X, y = Variable(data_adv, requires_grad=True), Variable(target)
            out = model(X)
            err_robust = (out.data.max(1)[1] != y.data).float().sum()
            robust_err_total += err_robust
            print(
                f'{1 - robust_err_total / (idx + 1):.2%} '
                f'({idx + 1 - int(robust_err_total)} / {idx + 1})')
    if not valid:
        print('not valid adversarial image')
    else:
        print('robust_err_total: ', robust_err_total * 1.0 / len(Y_data))


def main():
    # white-box attack
    # load model
    model = WideResNet().to(device)
    model.load_state_dict(torch.load(args.model_path, map_location=device))

    # load data
    X_adv_data = np.load(args.data_attak_path)
    X_data = np.load(args.data_path)
    Y_data = np.load(args.target_path)

    eval_adv_test_whitebox(model, device, X_adv_data, X_data, Y_data)


if __name__ == '__main__':
    main()