gmmclassifier_torch_attack_script_v3.py

# -*- coding: utf-8 -*-
"""GMMClassifier_Torch Attack Script v3.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1EQJf65dctV2ANzBvJUEmSaVr3XcXQKh4
"""

# !pip install foolbox

import time, torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import MNIST
import numpy as np
from foolbox import PyTorchModel, accuracy, samples, plot, criteria, distances
import foolbox.attacks as fa
import torchvision.models as models
import eagerpy as ep
# from imagenet_stubs.imagenet_2012_labels import label_to_name
import matplotlib.pyplot as plt

from sklearn.mixture import GaussianMixture
from sklearn.base import BaseEstimator, ClassifierMixin
from numpy.random import permutation

import torch.distributions as D

import foolbox as fb

from keras.datasets import mnist


if torch.cuda.is_available():
    device = 'cpu'
else:
    device = 'cpu'

class GMMSKlearnClassifier(BaseEstimator, ClassifierMixin):
    
    def __init__(self, n_components, n_features, covariance_type='spherical', max_iter=1000):
        self.n_component = n_components
        self.n_features = n_features
        self.covariance_type = covariance_type
        self.max_iter = max_iter
        
    def initialize_gmm(self, k, y, weights, means, Sigmas):
        
        self.y_classes_ = np.unique(y)

        self.model_list_ = np.asarray([GaussianMixture(self.n_component, 
                                                       covariance_type=self.covariance_type, 
                                                       max_iter=self.max_iter, 
                                                       means_init=means[i],
                                                       weights_init=weights[i],
                                                       precisions_init=Sigmas[i],
                                                       random_state=0) 
                                     for i in range(k)])
        
    def fit(self, X, y):
        self.y_classes_ = np.unique(y)
        
        # Create a Subset list. Each X Subset for each Y Class
        X_train_subsets = np.asarray([X[y==self.y_classes_[i]] for i in range(self.y_classes_.shape[0])])
        
        # Create a GMM list. Each GMM for each X Subset for each Y Class
        # Train the GMM model with the X subset
        # The model will learn and estimate the density distribution (Gaussian - mu and Sigma)
        self.model_list_ = np.asarray([GaussianMixture(self.n_component, 
                                                     covariance_type=self.covariance_type, 
                                                     max_iter=self.max_iter, 
                                                     random_state=0).fit(X_train_subsets[i]) 
                                     for i in range(self.y_classes_.shape[0])])
        
        # Prior list which is the X Subset ratio based on the Sample data
        self.logprior_list_ = np.asarray([
            np.log(X_train_subsets[i].shape[0] / X.shape[0]) for i in range(self.y_classes_.shape[0])
        ])
        
        return self
    
    def predict(self, X, method='mle'):
        
        result = None
        if method == 'mle':
            # use np.argmax to find the max for each N
            # predict_proba return (k x N), so axis=0
            result = self.y_classes_[np.argmax(self.predict_proba_likelihood(X), axis=0)]
        elif method == 'post':
            result = self.y_classes_[np.argmax(self.predict_proba_posterior(X), axis=0)]
        else:
            print(f'Unsupported classification method')
    
        return result
    
    def predict_proba_likelihood(self, X):
        # (k x N)
        loglikelihood_list = np.asarray([self.model_list_[i].score_samples(X) for i in range(self.y_classes_.shape[0])])       
        likelihood_list = np.exp(loglikelihood_list)
        
        # Sum the likelihood_list over the k dimension, keep the number of N
        # (k x N) / (1 x N)
        normalized = likelihood_list / likelihood_list.sum(0, keepdims=True)
        
        # return (k x N) where the sum of the probabilities for each N is 1
        return normalized
    
    def predict_proba_posterior(self, X):
        # (k x N)
        loglikelihood_list = np.asarray([self.model_list_[i].score_samples(X) for i in range(self.y_classes_.shape[0])])
        # Need to change logprior_list from (k, ) to (k, 1)
        # (k x N) + (k, 1)
        posterior_list = np.exp(loglikelihood_list + self.logprior_list_[:, np.newaxis])
        
        # Sum the posterior_list over the k dimension, keep the number of N
        # (k x N) / (1 x N)
        normalized = posterior_list / posterior_list.sum(0, keepdims=True)
        
        # return (k x N) where the sum of the probabilities for each N is 1
        return normalized
    
    def get_gmm_params(self):
        means_list = []
        weights_list = []
        covariances_list = []
        for i in range(self.y_classes_.shape[0]):
            means_list.append(self.model_list_[i].means_)
            weights_list.append(self.model_list_[i].weights_)
            covariances_list.append(self.model_list_[i].covariances_)
            
        return np.asarray(means_list), np.asarray(weights_list), np.asarray(covariances_list)


class GMMTorchClassifier(torch.nn.Module):
    
    def __init__(self, n_components, n_features, covariance_type='spherical', max_iter=1000):
        super().__init__()
        self.n_components = n_components
        self.n_features = n_features
        self.covariance_type = covariance_type
        self.max_iter = max_iter

    def fit(self, X, y):
        
        # 1. Use sklearn to train and get the means, weights, and variances parameters
        #
        X_np = X.detach().cpu().numpy()
        y_np = y.detach().cpu().numpy()
        
        sklearn_gmm_clf = GMMSKlearnClassifier(
            n_components=self.n_components,
            n_features=self.n_features,
            covariance_type=self.covariance_type,
            max_iter=self.max_iter,
        ).fit(X_np, y_np)
        
        means_list_np, weights_list_np, variances_list_np = sklearn_gmm_clf.get_gmm_params()
        
        self.y_classes_ = torch.unique(y).to(device=device)

        # Create a Subset list. Each X Subset for each Y Class
        X_train_subsets = [X[y==self.y_classes_[i]] for i in range(self.y_classes_.shape[0])]
        
        self.model_list_ = []
        for i in range(self.y_classes_.shape[0]):
            weights = torch.from_numpy(weights_list_np[i]).float().to(device=device)  # n_components
            means = torch.from_numpy(means_list_np[i]).float().to(device=device)      # n_components x n_features
            variances = torch.from_numpy(variances_list_np[i]).float().to(device=device) # n_components
            
            
            # covar = k x D x D 
            covariances_list = torch.empty((variances.shape[0], self.n_features, self.n_features)).to(device=device)
            for i in range(variances.shape[0]):
                covariances_list[i] = torch.eye(self.n_features).to(device=device) * variances[i]
    
            mix = D.Categorical(weights)
            comp = D.MultivariateNormal(means, scale_tril=torch.cholesky(covariances_list))
            gmm = D.MixtureSameFamily(mix, comp)
            self.model_list_.append(gmm)

    def forward(self, X):

        X = X.reshape(X.shape[0], -1)

        likelihood_list = self.predict_proba_likelihood(X)
        softmax = torch.nn.Softmax(dim=0)
        softmax_list = softmax(likelihood_list)

        return softmax_list.T

    def predict(self, X):
        
        likelihood_list = self.predict_proba_likelihood(X)
        softmax = torch.nn.Softmax(dim=0)
        softmax_list = softmax(likelihood_list)
        
        result = self.y_classes_[torch.argmax(softmax_list, axis=0)]
        return result

    def predict_proba_likelihood(self, X):
        # (k x N)
        loglikelihood_list = torch.empty((self.y_classes_.shape[0], X.shape[0])).to(device=device)
        for i in range(self.y_classes_.shape[0]):
            loglikelihood_list[i] = self.model_list_[i].log_prob(X)

        return loglikelihood_list
        

def load_data():
    (X_train, y_train), (X_test, y_test) = mnist.load_data()

    print('Training Data: {}'.format(X_train.shape))
    print('Training Labels: {}'.format(y_train.shape))
    print('Testing Data: {}'.format(X_test.shape))
    print('Testing Labels: {}'.format(y_test.shape))

    # preprocessing the images

    # convert each image to 1 dimensional array
    X_train = X_train.reshape(len(X_train),-1)
    X_test = X_test.reshape(len(X_test),-1)

    # normalize the data to 0 - 1
    X_train = X_train.astype(float) / 255.
    X_test = X_test.astype(float) / 255.

    y_classes = set(y_train)

    print(X_train.shape)
    print(X_train[0].shape)
    print(X_test.shape)
    print(X_test[0].shape)
    print(y_train.shape)
    print(y_test.shape)
    print(f'Y Class Labels: {y_classes}')

    perm_train = permutation(X_train.shape[0])
    perm_test = permutation(X_test.shape[0])

    # X_train_torch = torch.from_numpy(X_train[perm_train][:5000]).float().to(device=device)
    # y_train_torch = torch.from_numpy(y_train[perm_train][:5000]).float().to(device=device)
    X_train_torch = torch.from_numpy(X_train[perm_train][:]).float().to(device=device)
    y_train_torch = torch.from_numpy(y_train[perm_train][:]).float().to(device=device)
    X_test_torch = torch.from_numpy(X_test[perm_test][:100]).float().to(device=device)
    y_test_torch = torch.from_numpy(y_test[perm_test][:100]).float().to(device=device)

    print(X_train_torch.shape)
    print(y_train_torch.shape)
    print(X_train_torch.device)
    print(y_train_torch.device)

    return X_train_torch, y_train_torch, X_test_torch, y_test_torch, y_classes


def train_model(X_train_torch, y_train_torch):

    gmm_clf_torch = GMMTorchClassifier(n_components=60, n_features=X_train_torch.shape[1], covariance_type='spherical').to(device=device)
    gmm_clf_torch.fit(X_train_torch, y_train_torch)

    return gmm_clf_torch

def test_model(gmm_clf_torch, X_test_torch, y_test_torch):

    y_pred = gmm_clf_torch.predict(X_test_torch)
    acc = torch.sum(y_pred==y_test_torch)/y_test_torch.shape[0]
    print(f'Accuracy: {acc*100:.2f}%')    

def plot_images(images):
    fig, axs = plt.subplots(1, len(images), figsize=(24, 24))

    for i, ax in enumerate(axs.flat):
        image = images[i]
        image = np.squeeze(image, axis=0)
        ax.matshow(image)
        ax.axis('off')
        ax.set_title(f'Image {i}')

    fig.show()

def setup_fmodel(gmm_clf_torch):
    model = gmm_clf_torch.eval()
    fmodel = PyTorchModel(model, bounds=(0, 1))
    batch_size = 20
    images, labels = samples(fmodel, dataset="mnist", batchsize=20)

    clean_acc = accuracy(fmodel, images, labels)
    predictions = fmodel(images).argmax(-1)
    plot_images(images.cpu().numpy())
    print("labels     : ", [l for l in labels.cpu().numpy()])
    print("predictions: ", [l for l in predictions.cpu().numpy()])

    print(f"clean accuracy:  {clean_acc * 100:.1f} %")
    print("")    

    return fmodel

def main():


    X_train_torch, y_train_torch, X_test_torch, y_test_torch, y_classes = load_data()
    gmm_clf_torch = train_model(X_train_torch, y_train_torch)
    test_model(gmm_clf_torch, X_test_torch, y_test_torch)
    fmodel = setup_fmodel(gmm_clf_torch)

    test_size = 10  # assuming we are running 1000 tests for each attack
    batch_size = 2   # adjust your batch size according to memory limitation, as long as it's a divisor of 1000
    # test_size = 1000  # assuming we are running 1000 tests for each attack
    # batch_size = 50   # adjust your batch size according to memory limitation, as long as it's a divisor of 1000
    num_batch = int(test_size / batch_size)

    test_dataset = MNIST(root='./data', download=True, train=False, transform=transforms.ToTensor())
    test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

    Linf_epsilons = [
        0.1,
        0.2,
        0.3,
        0.4,
        0.5
    ]
    CW_epsilons = [
        0.1,
        1,
        10,
        100,
        1000
    ]
    L2_epsilons = [
        1,
        2,
        3,
        4,
        5
    ]
    attacks = [
        # gradient-based attacks
        # {"name": "Fast Gradient Method", "model": fa.L2FastGradientAttack(), "epsilons": L2_epsilons},
        # {"name": "Fast Gradient Sign Method", "model": fa.LinfFastGradientAttack(), "epsilons": Linf_epsilons},
        # {"name": "L2 Projected Gradient Descent", "model": fa.L2ProjectedGradientDescentAttack(), "epsilons": L2_epsilons},
        # {"name": "Linf Projected Gradient Descent", "model": fa.LinfProjectedGradientDescentAttack(), "epsilons": Linf_epsilons},
        # {"name": "Carlini & Wagner L2 Attack", "model": fa.L2CarliniWagnerAttack(), "epsilons": CW_epsilons},
        # # decision-based attacks
        # {"name": "Boundary Attack", "model": fa.BoundaryAttack(), "epsilons": L2_epsilons}, 
        {"name": "Additive Uniform Noise Attack", "model": fa.LinfAdditiveUniformNoiseAttack(), "epsilons": Linf_epsilons},
        # {"name": "Additive Gaussian Noise Attack", "model": fa.L2AdditiveGaussianNoiseAttack(), "epsilons": L2_epsilons},
        # {"name": "Linear Search Contrast Reduction Attack", "model": fa.LinearSearchContrastReductionAttack(distance=distances.linf), "epsilons": Linf_epsilons}, 
        # {"name": "Binary Search Contrast Reduction Attack", "model": fa.BinarySearchContrastReductionAttack(distance=distances.linf), "epsilons": Linf_epsilons}, 
        # {"name": "Gaussian Blur Attack", "model": fa.GaussianBlurAttack(distance=distances.linf), "epsilons": Linf_epsilons}
    ]

    fig, axs = plt.subplots(len(attacks))
    fig.tight_layout()

    attack_results = {}
    for attack in attacks:
        attack_results[attack['name']] = None

    for i, attack in enumerate(attacks):
        clean_accuracy = torch.zeros(num_batch)
        robust_accuracy_batch = torch.zeros([num_batch, len(CW_epsilons)])
        count = 0
        t1 = time.time()
        for images, labels in test_dataloader:
            images = images.to(device)
            labels = labels.to(device)
            if count == num_batch:
                break
            images = ep.astensor(images)
            labels = ep.astensor(labels)
            clean_accuracy[count] = fb.utils.accuracy(fmodel, images, labels)

            # The raw adversarial examples. This depends on the attack and we cannot make an guarantees about this output.
            # The clipped adversarial examples. These are guaranteed to not be perturbed more than epsilon and thus are the actual adversarial examples you want to visualize. Note that some of them might not actually switch the class. To know which samples are actually adversarial, you should look at the third tensor.
            # The third tensor contains a boolean for each sample, indicating which samples are true adversarials that are both misclassified and within the epsilon balls around the clean samples.
            raw_adv, clipped_adv, is_adv = attack["model"](fmodel, images, labels, epsilons=attack["epsilons"])

            success_ = is_adv.numpy()
            robust_accuracy_batch[count] = torch.from_numpy(1.0 - success_.mean(axis=-1))

            count += 1

        t2 = time.time()
        print(attack)
        print("Clean accuracy: {:.5f}".format(torch.mean(clean_accuracy)))
        print("Average time taken for each test: {} seconds".format((t2 - t1) / test_size))
        robust_accuracy = torch.mean(robust_accuracy_batch, dim=0)
        print("Accuracy with {} attack: {}".format(attack["name"], robust_accuracy))

        attack_results[attack['name']] = robust_accuracy

        # plot the robust accuracy against epsilons for the particular attack
        # axs[i].set_ylim([0, 1])
        # axs[i].plot(attack["epsilons"], robust_accuracy.numpy())
        # axs[i].set_title(attack["name"])
        # axs.set_ylim([0, 1])
        # axs.plot(attack["epsilons"], robust_accuracy.numpy())
        # axs.set_title(attack["name"])    

if __name__ == '__main__':
    main()