DefenseEvaluations.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# **************************************
# @Time    : 2018/9/17 1:23
# @Author  : Jiaxu Zou & Xiang Ling
# @Lab     : nesa.zju.edu.cn
# @File    : DefenseEvaluations.py
# **************************************

import argparse
import os
import random
import sys
import warnings

import numpy as np
import scipy.stats as jslib
import torch
import torch.nn.functional as F

sys.path.append('%s/../' % os.path.dirname(os.path.realpath(__file__)))
from RawModels.Utils.dataset import get_mnist_test_loader, get_cifar10_test_loader
from RawModels.ResNet import resnet20_cifar
from RawModels.MNISTConv import MNISTConvNet


# help functions
def pd_prediction(model, dataset, data_loader, epsilon, device):
    from Defenses.DefenseMethods.PD import PixelDefend
    pd = PixelDefend(model=model, defense_name='PD', dataset=dataset, pixel_cnn_dir='../Defenses/', device=device)
    model.eval()
    predicted_defended = []
    with torch.no_grad():
        for index, (images, labels) in enumerate(data_loader):
            images = images.detach().cpu().numpy()

            purified_images = pd.de_noising_samples(samples=images, batch_size=images.shape[0], eps=epsilon)

            rt_logits = model(torch.from_numpy(purified_images).to(device))
            rt_predicted = F.softmax(rt_logits, dim=1).cpu().numpy()
            predicted_defended.extend(rt_predicted)
        return np.array(predicted_defended)


def rt_prediction(model, dataset, data_loader, final_size, device):
    from Defenses.DefenseMethods.RT import RTDefense
    rt = RTDefense(model=model, defense_name='RT', dataset=dataset, device=device)
    model.eval()
    predicted_defended = []
    with torch.no_grad():
        for index, (images, labels) in enumerate(data_loader):
            transformed_images = rt.randomization_transformation(samples=images, original_size=images.shape[-1], final_size=final_size)
            transformed_images = transformed_images(device)

            rt_logits = model(transformed_images)
            rt_predicted = F.softmax(rt_logits, dim=1).cpu().numpy()
            predicted_defended.extend(rt_predicted)
    return np.array(predicted_defended)


def te_prediction(model, data_loader, level, device):
    from Defenses.DefenseMethods.TE import thermometer_encoding
    model.eval()
    predicted_defended = []
    with torch.no_grad():
        for index, (images, labels) in enumerate(data_loader):
            therm_inputs = thermometer_encoding(samples=images.to(device), level=level, device=device)
            therm_inputs = torch.from_numpy(therm_inputs).to(device)

            te_logits = model(therm_inputs)
            te_predicted = F.softmax(te_logits, dim=1).cpu().numpy()
            predicted_defended.extend(te_predicted)
    return np.array(predicted_defended)


def defense_utility_measure(pred_def, pred_raw, true_label):
    # compute the classification accuracy of raw model
    correct_prediction_raw = np.equal(np.argmax(pred_raw, axis=1), true_label)
    acc_raw = np.mean(correct_prediction_raw.astype(float))

    # compute the classification accuracy of defense-enhanced model
    correct_prediction_def = np.equal(np.argmax(pred_def, axis=1), true_label)
    acc_def = np.mean(correct_prediction_def.astype(float))

    # compute the Classification Accuracy Variance(CAV)
    cav_result = acc_def - acc_raw

    # find the index of correct predicted examples by defence-enhanced model and raw model
    idx_def = np.squeeze(np.argwhere(correct_prediction_def == True))
    idx_raw = np.squeeze(np.argwhere(correct_prediction_raw == True))
    idx = np.intersect1d(idx_def, idx_raw, assume_unique=True)

    # compute the Classification Rectify Ratio(CRR) & Classification Sacrifice Ratio(CSR)
    num_rectify = len(idx_def) - len(idx)
    crr_result = num_rectify / len(pred_def)

    num_sacrifice = len(idx_raw) - len(idx)
    csr_result = num_sacrifice / len(pred_def)

    # filter the correct prediction results
    pred_def_filter = pred_def[idx]
    pred_raw_filter = pred_raw[idx]

    # compute the Classification Confidence Variance(CCV)
    confidence_def = np.max(pred_def_filter, axis=1)
    confidence_raw = np.max(pred_raw_filter, axis=1)
    ccv_result = np.mean(np.absolute(confidence_def - confidence_raw))

    # compute the Classification Output Stability(COS)
    M = (pred_def_filter + pred_raw_filter) / 2.
    js_total = 0
    for i in range(len(M)):
        js = 0.5 * jslib.entropy(pred_def_filter[i], M[i]) + 0.5 * jslib.entropy(pred_raw_filter[i], M[i])
        js_total += js
    cos_result = js_total / len(M)

    return acc_raw, acc_def, cav_result, crr_result, csr_result, ccv_result, cos_result


def prediction(model, test_loader, device):
    print('\nThe #{}# model is evaluated on the testing dataset loader ...'.format(model.model_name))

    model = model.to(device)
    model.eval()

    prediction = []
    true_labels = []
    with torch.no_grad():
        for images, labels in test_loader:
            images = images.to(device)
            labels = labels.to(device)
            logits = model(images)
            predicted = F.softmax(logits, dim=1).cpu().numpy()

            prediction.extend(predicted)
            true_labels.extend(labels)
    prediction = np.array(prediction)
    true_labels = np.array(true_labels)
    return prediction, true_labels


def main(args):
    # Device configuration
    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_index
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # Set the random seed manually for reproducibility.
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    torch.manual_seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(args.seed)
    np.random.seed(args.seed)
    random.seed(args.seed)

    dataset = args.dataset.upper()
    assert dataset == 'MNIST' or dataset == 'CIFAR10'

    # load the raw model / testing dataset loader
    raw_model_location = '{}{}/model/{}_raw.pt'.format('../RawModels/', dataset, dataset)
    if dataset == 'MNIST':
        raw_model = MNISTConvNet().to(device)
        raw_model.load(path=raw_model_location, device=device)
        test_loader = get_mnist_test_loader(dir_name='../RawModels/MNIST/', batch_size=30)
    else:
        raw_model = resnet20_cifar().to(device)
        raw_model.load(path=raw_model_location, device=device)
        test_loader = get_cifar10_test_loader(dir_name='../RawModels/CIFAR10/', batch_size=25)
    raw_model.eval()

    # get predictions of the raw model on test datasets
    predicted_raw, true_label = prediction(model=raw_model, test_loader=test_loader, device=device)

    re_train_defenses = {'NAT', 'EAT', 'PAT', 'DD', 'IGR'}
    input_transformation_defenses = {'EIT', 'RT', 'PD', 'TE'}
    other_defenses = {'RC'}

    defense_name = args.defense.upper().strip()
    if defense_name in re_train_defenses:
        print('\nthe ##{}## defense is a kind of complete defenses that retrain the model'.format(defense_name))
        # load the defense-enhanced model
        defended_model_location = '{}/{}/{}_{}_enhanced.pt'.format('../DefenseEnhancedModels', defense_name, dataset, defense_name)
        defended_model = MNISTConvNet().to(device) if dataset == 'MNIST' else resnet20_cifar().to(device)
        defended_model.load(path=defended_model_location, device=device)
        defended_model.eval()
        predicted_defended, _ = prediction(model=defended_model, test_loader=test_loader, device=device)
        raw_acc, def_acc, cav, crr, csr, ccv, cos = defense_utility_measure(predicted_defended, predicted_raw, true_label)

    elif defense_name in input_transformation_defenses:
        print('\nthe ##{}## defense is a kind of complete defense hat need to transform the images ... '.format(defense_name))
        if defense_name == 'EIT':

            from Defenses.DefenseMethods.EIT import EITDefense, TransformedDataset
            eit_params = {
                'crop_size': args.crop_size,
                'lambda_tv': args.lambda_tv,
                'JPEG_quality': args.JPEG_quality,
                'bit_depth': args.bit_depth
            }
            defended_model = MNISTConvNet().to(device) if dataset == 'MNIST' else resnet20_cifar().to(device)
            EIT = EITDefense(model=defended_model, defense_name=defense_name, dataset=dataset, re_training=False, training_parameters=None,
                             device=device, **eit_params)
            transformed_test_data_numpy, test_label_numpy = EIT.transforming_dataset(data_loader=test_loader)
            transformed_test_dataset = TransformedDataset(images=torch.from_numpy(transformed_test_data_numpy),
                                                          labels=torch.from_numpy(test_label_numpy), dataset=dataset, transform=None)
            transformed_test_loader = torch.utils.data.DataLoader(transformed_test_dataset, batch_size=100, shuffle=False)

            defended_model_location = '{}/{}/{}_{}_enhanced.pt'.format('../DefenseEnhancedModels', defense_name, dataset, defense_name)
            defended_model = defended_model.to(device)
            defended_model.load(path=defended_model_location, device=device)
            defended_model.eval()

            predicted_defended, _ = prediction(model=defended_model, test_loader=transformed_test_loader, device=device)
            raw_acc, def_acc, cav, crr, csr, ccv, cos = defense_utility_measure(predicted_defended, predicted_raw, true_label)

        elif defense_name == 'RT':
            final_size = args.resize
            assert isinstance(final_size, int)
            warnings.warn(message='For the RT defense, the #resize# parameter is specified as {}, please check ...'.format(final_size))
            predicted_defended = rt_prediction(model=raw_model, dataset=dataset, data_loader=test_loader, final_size=final_size, device=device)
            # test the utility performance of defended model
            raw_acc, def_acc, cav, crr, csr, ccv, cos = defense_utility_measure(predicted_defended, predicted_raw, true_label)
        elif defense_name == 'PD':
            epsilon = args.epsilon
            warnings.warn(message='For the PixelDefend defense, the #epsilon# parameter is specified as {}, please check ...'.format(epsilon))
            predicted_defended = pd_prediction(model=raw_model, dataset=dataset, data_loader=test_loader, epsilon=epsilon, device=device)
            raw_acc, def_acc, cav, crr, csr, ccv, cos = defense_utility_measure(predicted_defended, predicted_raw, true_label)
        else:
            level = args.level
            assert defense_name == 'TE' and isinstance(level, int)
            warnings.warn(message='For the TE defense, the #level# parameter is specified as {}, please check ...'.format(level))

            # load the defense-enhanced model (for TE)
            defended_model_location = '{}/{}/{}_{}_enhanced.pt'.format('../DefenseEnhancedModels', defense_name, dataset, defense_name)
            defended_model = MNISTConvNet(thermometer=True, level=level) if dataset == 'MNIST' else resnet20_cifar(thermometer=True, level=level)
            defended_model = defended_model.to(device)
            defended_model.load(path=defended_model_location, device=device)
            defended_model.eval()
            predicted_defended = te_prediction(model=defended_model, data_loader=test_loader, level=level, device=device)
            raw_acc, def_acc, cav, crr, csr, ccv, cos = defense_utility_measure(predicted_defended, predicted_raw, true_label)

    else:
        if defense_name == 'RC':
            print('\n##{}## defense is a kind of region-based classification defenses ... '.format(defense_name))
            from Defenses.DefenseMethods.RC import RCDefense
            num_points = 1000
            radius = args.radius
            rc = RCDefense(model=raw_model, defense_name='RC', dataset=dataset, device=device, num_points=num_points)

            predicted_defended = []
            with torch.no_grad():
                for index, (images, labels) in enumerate(test_loader):
                    rc_labels = rc.region_based_classification(samples=images, radius=radius)
                    predicted_defended.extend(rc_labels)
            predicted_defended = np.array(predicted_defended)

            # classification accuracy of defense-enhanced model
            correct_prediction_def = np.equal(predicted_defended, true_label)
            def_acc = np.mean(correct_prediction_def.astype(float))
            # classification accuracy of raw model
            correct_prediction_raw = np.equal(np.argmax(predicted_raw, axis=1), true_label)
            raw_acc = np.mean(correct_prediction_raw.astype(float))
            # Classification Accuracy Variance(CAV)
            cav = def_acc - raw_acc

            # Find the index of correct predicted examples by defence-enhanced model and raw model
            idx_def = np.squeeze(np.argwhere(correct_prediction_def == True))
            idx_raw = np.squeeze(np.argwhere(correct_prediction_raw == True))
            idx = np.intersect1d(idx_def, idx_raw, assume_unique=True)

            # Compute the Classification Rectify Ratio(CRR) & Classification Sacrifice Ratio(CSR)
            crr = (len(idx_def) - len(idx)) / len(predicted_raw)
            csr = (len(idx_raw) - len(idx)) / len(predicted_raw)
            ccv = cos = 0

        else:
            raise ValueError('{} is not supported!!!'.format(defense_name))

    print("****************************")
    print("The utility evaluation results of the {} defense for {} Dataset are as follow:".format(defense_name, dataset))
    print('Acc of Raw Model:\t\t{:.2f}%'.format(raw_acc * 100))
    print('Acc of {}-enhanced Model:\t{:.2f}%'.format(defense_name, def_acc * 100))
    print('CAV: {:.2f}%'.format(cav * 100))
    print('CRR: {:.2f}%'.format(crr * 100))
    print('CSR: {:.2f}%'.format(csr * 100))
    print('CCV: {:.2f}%'.format(ccv * 100))
    print('COS: {:.4f}'.format(cos))
    print("****************************")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--dataset', type=str, default='CIFAR10', help='the dataset should be MNIST or CIFAR10')
    parser.add_argument('--gpu_index', type=str, default='0', help="gpu index to use")
    parser.add_argument('--seed', type=int, default=100, help='the default random seed for numpy and torch')

    parser.add_argument('--defense', type=str, default='PD', help="the defense method to be evaluated")

    # default argument for the EIT defense (only crop_size is different)
    parser.add_argument('--crop_size', type=int, default=30, help='the cropping size (26 for mnist and 30 for cifar10)')
    parser.add_argument('--bit_depth', type=int, default=4, help='the quantization level of pixel value')
    parser.add_argument('--JPEG_quality', type=int, default=85, help='the JPEG quality to compress with')
    parser.add_argument('--lambda_tv', type=float, default=0.03, help='the total variance minimization weight')

    # default argument for the RT defense in MNIST (final_size=31) or in CIFAR10 (final_size=36)
    parser.add_argument('--resize', type=int, default=36, help='the final size parameter only for the RT defense')
    # default argument for the PD defense in MNIST (epsilon=76.5) or in CIFAR10 (epsilon=16)
    parser.add_argument('--epsilon', type=float, default=0.0627, help="radius of e_ball")
    # default argument for the TE defense (level=16)
    parser.add_argument('--level', type=int, default=16, help='the discretization level of pixel value')
    # default argument for the RC defense in MNIST (radius=0.3) or in CIFAR10 (radius=0.02)
    parser.add_argument('--radius', type=float, default=0.02, help='in the case of not search radius r, we set the radius of the hypercube')

    arguments = parser.parse_args()
    main(arguments)