<repo_name>ChengYuHsieh/explanation_robustness<file_sep>/bbmp.py
import torch
from torch.autograd import Variable
from torchvision import models
import cv2
import sys
import numpy as np


def tv_norm(input, tv_beta):
    img = input[0, 0, :, :]
    row_grad = torch.mean(torch.abs((img[:-1 , :] - img[1 :, :])).pow(tv_beta))
    col_grad = torch.mean(torch.abs((img[: , :-1] - img[: , 1 :])).pow(tv_beta))
    return row_grad + col_grad

def bbmp(X, y, model, device):
    print('Computing explanation by BBMP...')
    tv_beta = 3
    learning_rate = 0.1
    max_iterations = 1000
    l1_coeff = 10
    tv_coeff = 0.01

    img = X.cpu().numpy().reshape(28, 28)
    blurred_img = cv2.GaussianBlur(img, (11, 11), 5)

    mask_init = np.ones((28, 28), dtype = np.float32)
    mask = torch.Tensor(mask_init).to(device).view(1, 1, 28, 28)
    mask.requires_grad = True

    optimizer = torch.optim.Adam([mask], lr=learning_rate)
    img = torch.Tensor(img).to(device).view(1, 1, 28, 28)
    blurred_img = torch.Tensor(blurred_img).to(device).view(1, 1, 28, 28)
    category = y.item()

    for i in range(max_iterations):
        perturbated_input = img.mul(mask) + \
                            blurred_img.mul(1-mask)

        noise = torch.randn(1, 1, 28, 28).to(device)
        
        perturbated_input = perturbated_input + noise

        outputs = model(perturbated_input)
        loss = l1_coeff*torch.mean(mask) + \
                tv_coeff*tv_norm(mask, tv_beta) + outputs[0, category]

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # Optional: clamping seems to give better results
        mask.data.clamp_(0, 1)
    
    return mask.data.cpu().numpy()<file_sep>/requirements.txt
absl-py==0.7.1
astor==0.8.0
astroid==2.3.1
attrs==19.1.0
backcall==0.1.0
bleach==3.1.0
certifi==2019.6.16
cffi==1.5.2
chardet==2.3.0
cleverhans==3.0.1
cloudpickle==1.2.2
command-not-found==0.3
cycler==0.10.0
dask==0.20.0
decorator==4.3.0
-e git+https://github.com/marcoancona/DeepExplain.git@<PASSWORD>#egg=deepexplain
defusedxml==0.6.0
entrypoints==0.3
fail2ban==0.9.3
future==0.17.1
gast==0.2.2
google-pasta==0.1.7
grpcio==1.21.1
gym==0.14.0
h5py==2.9.0
ipykernel==5.1.1
ipython==7.5.0
ipython-genutils==0.2.0
ipywidgets==7.4.2
isort==4.3.21
jedi==0.14.0
Jinja2==2.10.1
joblib==0.9.4
jsonschema==3.0.1
jupyter==1.0.0
jupyter-client==5.2.4
jupyter-console==6.0.0
jupyter-core==4.5.0
kaggle==1.5.5
Keras==2.2.4
Keras-Applications==1.0.8
Keras-Preprocessing==1.1.0
kiwisolver==1.0.1
language-selector==0.1
lazy-object-proxy==1.4.1
llvmlite==0.29.0
Markdown==3.1.1
MarkupSafe==1.1.1
matplotlib==3.0.1
mccabe==0.6.1
mistune==0.8.4
mnist==0.2.2
nbconvert==5.5.0
nbformat==4.4.0
networkx==2.2
nose==1.3.7
notebook==5.7.8
numba==0.44.1
numpy==1.15.3
olympic==0.1.5
opencv-python==4.1.1.26
pandas==0.24.2
pandocfilters==1.4.2
parso==0.5.0
pexpect==4.7.0
pickleshare==0.7.5
Pillow==5.3.0
ply==3.7
prometheus-client==0.7.1
prompt-toolkit==2.0.9
protobuf==3.8.0
ptyprocess==0.6.0
pycodestyle==2.5.0
pycparser==2.14
pycurl==7.43.0
pyglet==1.3.2
Pygments==2.4.2
pygobject==3.20.0
pyinotify==0.9.6
pylint==2.4.2
pyparsing==2.2.2
pyrsistent==0.15.2
python-apt==1.1.0b1+ubuntu0.16.4.8
python-dateutil==2.7.5
python-debian==0.1.27
python-slugify==3.0.3
python-systemd==231
pytz==2019.1
PyWavelets==1.0.1
PyYAML==5.1.1
pyzmq==18.0.1
qtconsole==4.5.1
requests==2.9.1
scikit-image==0.14.1
scikit-learn==0.17
scipy==1.1.0
screen-resolution-extra==0.0.0
Send2Trash==1.5.0
simplejson==3.8.1
six==1.12.0
ssh-import-id==5.5
tensorboard==1.14.0
tensorboardX==1.8
tensorflow-estimator==1.14.0
tensorflow-gpu==1.14.0
tensorflow-probability==0.7.0
termcolor==1.1.0
terminado==0.8.2
testpath==0.4.2
text-unidecode==1.2
toolz==0.9.0
torch==1.2.0
torchfile==0.1.0
torchvision==0.4.0
tornado==6.0.3
tqdm==4.34.0
traitlets==4.3.2
typed-ast==1.4.0
ufw==0.35
unattended-upgrades==0.1
urllib3==1.24.3
virtualenv==16.6.1
visdom==0.1.8.8
wcwidth==0.1.7
webencodings==0.5.1
websocket-client==0.56.0
Werkzeug==0.15.4
widgetsnbextension==3.4.2
wrapt==1.11.2
xkit==0.0.0
<file_sep>/pgd.py
import torch
from torch.autograd import Variable
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.transforms.functional as F
import numpy as np


def sample_eps_L2(num_dim, epsilon, N, anchor_masks, device):
    if num_dim == 0:
        return torch.zeros_like(anchor_masks).float()
    normal_deviates = np.random.normal(size=(N, num_dim))
    total_dim = normal_deviates.shape[1]
    radius = np.linalg.norm(normal_deviates, axis=1, keepdims=True)
    normal_deviates = normal_deviates * (np.random.rand(N, total_dim) ** (1.0 / total_dim))
    points = normal_deviates * (epsilon) / radius
    samples = torch.zeros_like(anchor_masks).float()
    samples[~anchor_masks] = torch.FloatTensor(points).to(device).view(-1)
    return samples


def pgd(Xs, ys, model, epsilons, norm, niters, step_size, anchor_masks, targeted, targets=None,
        box_min=0., box_max=1., verbose=False, multi_start=False, device=None): 
    if multi_start: # currently only support only example at a time
        assert Xs.shape == (1, 1, 28, 28)
        assert ys.shape == (1,)
        assert epsilons.shape == (1,)
        num_start = 1000
        anchor_masks = anchor_masks.repeat(num_start, 1, 1, 1)
        epsilons = epsilons.repeat(num_start)
        num_free_dim = 28*28 - anchor_masks[0].sum()
        if norm == 2:
            samples = sample_eps_L2(num_free_dim, epsilons[0].item(), num_start, anchor_masks, device)
            Xs_pgd = Variable(Xs.data + samples, requires_grad=True)
            Xs_pgd.data = torch.clamp(Xs_pgd.data, 0.0, 1.0)
        else:
            raise ValueError('norm not supported')

        if targeted:
            targets = targets.repeat(num_start).view(-1, 1)
        else:
            ys = ys.repeat(num_start).view(-1, 1)
    else:
        Xs_pgd = Variable(Xs.data, requires_grad=True)

        if targeted:
            targets = targets.view(-1, 1)
        else:
            ys = ys.view(-1, 1)


    for i in range(niters):
        outputs = model(Xs_pgd)
        if targeted:
            loss = (torch.gather(outputs, 1, targets)).sum() 
        else:
            loss = -((torch.gather(outputs, 1, ys)).sum())
    
        loss.backward()
    
        cur_grad = Xs_pgd.grad.data
        
        eta = cur_grad
        eta[anchor_masks] = 0. 
        zero_mask = eta.view(eta.shape[0], -1).norm(norm, dim=-1) == 0.
        eta = (eta * step_size /
               eta.view(eta.shape[0], -1).norm(norm, dim=-1).view(-1, 1, 1, 1))
        eta[zero_mask] = 0.

        Xs_pgd = Variable(Xs_pgd.data + eta, requires_grad=True)
        Xs_pgd.data = torch.clamp(Xs_pgd.data, 0.0, 1.0) 
    
        eta = Xs_pgd.data - Xs.data
        
        mask = eta.view(eta.shape[0], -1).norm(norm, dim=1) <= epsilons
        scaling_factor = eta.view(eta.shape[0], -1).norm(norm, dim=1)
        scaling_factor[mask] = epsilons[mask]
        eta *= (epsilons / scaling_factor).view(-1, 1, 1, 1)

        Xs_pgd.data = Xs.data + eta

    ys_pred_pgd = model(Xs_pgd)
    if targeted:
        successes = ys_pred_pgd.data.max(1)[1] == targets.data.view(-1)
    else:
        successes = ys_pred_pgd.data.max(1)[1] != ys.data.view(-1)
    
    # print('attack success on {} of {}'.format(successes.sum().item(), len(Xs_pgd)))

    if multi_start:
        success = successes.any(dim=0, keepdim=True)
        success_Xs_pgd = Xs_pgd.data[successes][[0]] if success.item() else Xs_pgd.data[[0]]
        return success, success_Xs_pgd
    else:
        return successes, Xs_pgd.data
    
            
    
    <file_sep>/README.txt
The main file is "mnist_example.ipynb". Please use Jupyter notebook to open the file.
We also provide an HTML file "mnist_example.html" containing the executed results of the notebook. The file could be open in standard web browsers.

The code is written in Python 3.
- Install python dependencies:
1. `pip install -r requirements.txt`

Reference:
1. https://github.com/PAIR-code/saliency
2. https://github.com/marcoancona/DeepExplain
3. https://github.com/ruthcfong/pytorch-explain-black-box
4. https://github.com/chihkuanyeh/saliency_evaluation<file_sep>/expl_utils.py
import os
import sys
import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
from pgd import pgd
from sklearn import linear_model


def binary_search(cond, Xs, ys, epsilons, anchor_masks, targeted, targets,
                  upper=np.inf, lower=0.0, tol=0.00001, max_steps=100, upscale_mult=10.0,
                  downscale_mult=10.0, lower_limit=1e-8, upper_limit=1e+3):
    """batch binary search for attack approximation of robustness
    """
    batch_size = len(Xs)
    upper_not_found = np.full(batch_size, True)
    lower_not_found = np.full(batch_size, True)
    upper = np.full(batch_size, upper)
    lower = np.full(batch_size, lower)
    
    Xs_pgd = torch.zeros_like(Xs)  # placeholder for perturbed images
    
    step = 0
    run_idx = np.arange(batch_size)
    while len(run_idx) != 0:
        if targeted:
            cur_successes, cur_Xs_pgd = cond(Xs[run_idx], ys[run_idx], epsilons[run_idx], anchor_masks[run_idx],
                                             targeted, targets[run_idx])
        else:
            cur_successes, cur_Xs_pgd = cond(Xs[run_idx], ys[run_idx], epsilons[run_idx], anchor_masks[run_idx],
                                             targeted, None) 
        for cur_i, run_i in enumerate(run_idx):
            if cur_successes[cur_i]: # success at current value (attack success)
                Xs_pgd[run_i] = cur_Xs_pgd[cur_i]
                if lower_not_found[run_i]: # always success, we have not found the true lower bound
                    upper[run_i] = epsilons[run_i]
                    epsilons[run_i] /= downscale_mult  # downscale search range, try to find an pseudo lower bound
                else:
                    upper[run_i] = epsilons[run_i]
                    epsilons[run_i] = 0.5 * (lower[run_i] + upper[run_i])
                upper_not_found[run_i] = False # when initial is a success we always have a valid upper bound
            else:
                if upper_not_found[run_i]:
                    lower[run_i] = epsilons[run_i]
                    epsilons[run_i] *= upscale_mult
                else:
                    lower[run_i] = epsilons[run_i]
                    epsilons[run_i] = 0.5 * (lower[run_i] + upper[run_i]) # when initial is a failure we have a pseudo lower bound
                lower_not_found[run_i] = False
        step += 1
        not_meet_tolerance = (upper - lower) > tol  # not meet tolerance
        not_exceed_limit = np.logical_and((epsilons < upper_limit), (epsilons > lower_limit)) # not exceed limit
        run_cond = np.logical_and(not_meet_tolerance, not_exceed_limit)
        run_idx = np.arange(batch_size)[run_cond] 
        if step >= max_steps:
            break

    return upper, Xs_pgd


def uniform_sampling(X, num_samples, cur_anchor=None):
    """direct uniform sampling
    """
    if cur_anchor is not None:
        # sample from the rest of un-anchored features
        samples = np.repeat(cur_anchor.reshape(1, -1), num_samples, axis=0)
        samples_flatten = np.random.choice([0, 1], size=int(num_samples * (28*28 - cur_anchor.sum())), replace=True, p=[0.8, 0.2])
        samples[samples!=1] = samples_flatten
    else:
        samples = np.random.choice([0,1], size=(num_samples, X.shape[1], X.shape[2], X.shape[3]), replace=True, p=[0.8, 0.2])
    return samples, samples


def iterative_sampling(X, num_samples, cur_anchor, num_new_anchor):
    samples = np.repeat(cur_anchor.reshape(1, -1), num_samples, axis=0)
    samples_flatten = np.zeros((num_samples, int(28*28-cur_anchor.sum())))
    samples_flatten[:, :num_new_anchor] = 1
    samples_flatten = np.apply_along_axis(np.random.permutation, 1, samples_flatten)
    samples[samples == 0] = samples_flatten.reshape(-1)
    Zs = samples_flatten
    return samples, Zs


def evaluate_robustness(Xs, ys, model, targeted, targets, norm, anchors, eps_start, step_size, niters, device,
                        reverse_anchor=False, batch_size=10000):
    """evaluating robustness for different anchors for a given X"""
    assert Xs.shape == (1, 1, 28, 28)
    assert ys.shape == (1,)
    if targeted:
        assert targets.shape == (1,)
    
    total = range(len(anchors))
    num_batch = len(total) // batch_size
    if len(total) % batch_size != 0:
        num_batch += 1
    batches = [total[i*batch_size:(i+1)*batch_size] for i in range(num_batch)]
    robust_ub = np.full(len(total), np.inf)
    for batch in batches:
        Xs_repeat = Xs.repeat(len(batch), 1, 1, 1)
        ys_repeat = ys.repeat(len(batch))
        eps_repeat = np.full(len(batch), eps_start)
        if targeted:
            targets_repeat = targets.repeat(len(batch))
        cur_anchors = torch.BoolTensor(anchors[batch]).to(device)
        if reverse_anchor:
            cur_anchors = ~cur_anchors
        def binary_search_cond(Xs, ys, epsilons, anchor_masks, targeted, targets):
            epsilons = torch.FloatTensor(epsilons).to(device)
            successes, Xs_pgd = pgd(Xs, ys, model, epsilons, norm, niters, step_size, anchor_masks, targeted, targets=targets,
                                    box_min=0., box_max=1., verbose=False, multi_start=False, device=device)
            return successes, Xs_pgd
        if targeted:
            robust_ub[batch], _ = binary_search(binary_search_cond, Xs_repeat, ys_repeat, eps_repeat, cur_anchors, targeted, targets_repeat, tol=0.001, max_steps=25)
        else:
            robust_ub[batch], _ = binary_search(binary_search_cond, Xs_repeat, ys_repeat, eps_repeat, cur_anchors, targeted, None, tol=0.001, max_steps=25)
    
    return robust_ub


def greedy(X, y, model, targeted, target, norm, eps_start, step_size, niters, device, reverse_anchor, percentage=0.3, writer=None):
    anchor_map = np.zeros(28 * 28)
    X_np = X.cpu().numpy()
    cur_anchor = np.zeros(28 * 28)
    cur_val = 1000  # for greedy ranking
    while cur_anchor.sum() < 784 * percentage:
        n_sample = int(784 - cur_anchor.sum())
        Zs = np.eye(n_sample)
        anchor_samples = np.repeat(cur_anchor.reshape(1, -1), n_sample, axis=0)
        anchor_samples[anchor_samples != 1] = Zs.reshape(-1)
        anchor_samples = anchor_samples.reshape(n_sample, 1, 28, 28)

        robust_ub = evaluate_robustness(X, y, model, targeted, target, norm, anchor_samples, eps_start,
                                        step_size, niters, device, reverse_anchor=reverse_anchor)
        robust_ub[robust_ub == np.inf] = np.linalg.norm(X_np - np.zeros_like(X_np))

        ks = np.ones(n_sample)
        feature_robust = kernel_regression(Zs, ks, robust_ub)
        if reverse_anchor:
            max_idx = feature_robust.argsort()[:int(784*0.05)]
        else:
            max_idx = feature_robust.argsort()[-int(784*0.05):]  # select four features at a single greedy step (speed up)
        max_idx = np.where(cur_anchor != 1)[0][max_idx]
        cur_anchor[max_idx] = 1
        anchor_map[max_idx] = cur_val
        cur_val -= 2

    anchor_map = anchor_map.reshape(1, 1, 28, 28)
    return anchor_map


def one_step_banzhaf(X, y, model, targeted, target, norm, eps_start, step_size, niters, device, reverse_anchor, writer=None):
    n_sample = 10000

    # anchor_samples = uniform_sampling(X, num_samples)
    X_np = X.cpu().numpy()
    Xs = np.repeat(X_np.reshape(1, -1), n_sample, axis=0)
    Zs = np.apply_along_axis(shap_sampling, 1, Xs, p=0.5)
    anchor_samples = Zs.reshape(n_sample, 1, 28, 28)
    
    robust_ub = evaluate_robustness(X, y, model, targeted, target, norm, anchor_samples, eps_start,
                                    step_size, niters, device, reverse_anchor=reverse_anchor)
    robust_ub[robust_ub == np.inf] = np.linalg.norm(X_np - np.zeros_like(X_np))

    ks = np.ones(n_sample)
    feature_robust = kernel_regression(Zs, ks, robust_ub)
    
    if reverse_anchor:
        feature_robust = -feature_robust
        
    if writer is None:
        return feature_robust
    else:
        raise NotImplementedError()
        fig, name = vis(X, feature_robust, np.zeros_like(X), 'ls_regression', None, tensorboard=True, rescale=False)
        writer.add_figure(name, fig)


def forward_batch(model, input, batchsize):
    inputsize = input.shape[0]
    for count in range((inputsize - 1) // batchsize + 1):
        end = min(inputsize, (count + 1) * batchsize)
        if count == 0:
            tempinput = input[count * batchsize:end]
            out = model(tempinput.cuda())
            out = out.data.cpu().numpy()
        else:
            tempinput = input[count * batchsize:end]
            temp = model(tempinput.cuda()).data
            out = np.concatenate([out, temp.cpu().numpy()], axis=0)
    return out


def shap(X, label, pdt, model, n_sample):
    X = X.cpu().numpy()
    Xs = np.repeat(X.reshape(1, -1), n_sample, axis=0)
    Xs_img = Xs.reshape(n_sample, 1, 28, 28)

    Zs = np.apply_along_axis(sample_shap_Z, 1, Xs)
    Zs_real = np.copy(Zs)
    Zs_real[Zs == 1] = Xs[Zs == 1]
    Zs_real_img = Zs_real.reshape(n_sample, 1, 28, 28)
    Zs_img = Variable(torch.tensor(Xs_img - Zs_real_img), requires_grad=False).float()
    out = forward_batch(model, Zs_img, 5000)
    ys = out[:, label]

    ys = pdt.data.cpu().numpy() - ys
    ks = np.apply_along_axis(shap_kernel, 1, Zs, X=X.reshape(-1))

    expl = kernel_regression(Zs, ks, ys)

    return expl


def kernel_regression(Is, ks, ys):
    """
    *Inputs:
        I: sample of perturbation of interest, shape = (n_sample, n_feature)
        K: kernel weight
    *Return:
        expl: explanation minimizing the weighted least square
    """
    n_sample, n_feature = Is.shape
    IIk = np.matmul(np.matmul(Is.transpose(), np.diag(ks)), Is)
    Iyk = np.matmul(np.matmul(Is.transpose(), np.diag(ks)), ys)
    expl = np.matmul(np.linalg.pinv(IIk), Iyk)
    return expl


def nCr(n, r):
    f = math.factorial
    return f(n) // f(r) // f(n-r)


def shap_sampling(X, p=0.5):
    nz_ind = np.nonzero(X)[0]
    nz_ind = np.arange(X.shape[0])
    num_nz = len(nz_ind)
    bb = 0
    while bb == 0 or bb == num_nz:
        aa = np.random.rand(num_nz)
        bb = np.sum(aa > p)
    sample_ind = np.where(aa > p)
    Z = np.zeros(len(X))
    Z[nz_ind[sample_ind]] = 1
    return Z


def sample_shap_Z(X):
    nz_ind = np.nonzero(X)[0]
    nz_ind = np.arange(X.shape[0])
    num_nz = len(nz_ind)
    bb = 0
    while bb == 0 or bb == num_nz:
        aa = np.random.rand(num_nz)
        bb = np.sum(aa > 0.5)
    sample_ind = np.where(aa > 0.5)
    Z = np.zeros(len(X))
    Z[nz_ind[sample_ind]] = 1

    return Z


def shap_kernel(Z, X):
    M = X.shape[0]
    z_ = np.count_nonzero(Z)
    return (M-1) * 1.0 / (z_ * (M - 1 - z_) * nCr(M - 1, z_))


def greedy_as(X, y, model, targeted, target, norm, eps_start, step_size, niters, device, reverse_anchor, percentage=0.5):
    print('Computing explanation by Greedy-AS...')
    anchor_map = np.zeros(28 * 28)
    X_np = X.cpu().numpy()
    cur_anchor = np.zeros(28 * 28)
    cur_val = 1000.  # for greedy ranking
    
    while cur_anchor.sum() < 784 * percentage:
        print('number of relevant features selected:', cur_anchor.sum())
        n_sample = 5000
        anchor_samples, Zs = uniform_sampling(X_np, n_sample, cur_anchor)
        anchor_samples = anchor_samples.reshape(n_sample, 1, 28, 28)
        robust_ub = evaluate_robustness(X, y, model, targeted, target, norm, anchor_samples, eps_start,
                                        step_size, niters, device, reverse_anchor=reverse_anchor)
        robust_ub[robust_ub == np.inf] = np.linalg.norm(X_np - np.zeros_like(X_np))

        ks = np.ones(n_sample)
        feature_robust = kernel_regression(Zs, ks, robust_ub)
        if reverse_anchor:
            max_idx = feature_robust[cur_anchor != 1].argsort()[:int(784*0.05)]
        else:
            max_idx = feature_robust[cur_anchor != 1].argsort()[-int(784*0.05):]
        max_idx = np.where(cur_anchor != 1)[0][max_idx]
        cur_anchor[max_idx] = 1
        anchor_map[max_idx] = cur_val
        cur_val -= 2

    anchor_map = anchor_map.reshape(1, 1, 28, 28)
    return anchor_map


def plot_eval_curve(X, y, model, target, norm, eps_start, step_size, niters, writer, anchor, method, reverse_anchor):
    plot_step = 0
    for i in range(5, 50, 5):
        threshold = np.percentile(anchor, 100-i)
        cur_anchor = np.copy(anchor)
        cur_anchor[anchor <= threshold] = 0
        cur_anchor[anchor > threshold] = 1
        if reverse_anchor:
            cur_anchor = np.ones_like(cur_anchor) - cur_anchor
        def binary_search_cond(current_epsilons, cur_anchor_masks):
            success, X_pgd, margin = pgd(None, X.repeat(len(current_epsilons), 1, 1, 1), y.repeat(len(current_epsilons)), 
                        [model], current_epsilons, norm, niters=niters, step_size=step_size, alpha=[1],
                        anchor_idx=cur_anchor_masks, target=target.repeat(len(current_epsilons)), multi_start=True, multi_success=False)
            return success, X_pgd, margin
        robust_ub, X_pgd, _ = binary_search(binary_search_cond, np.full((1, 1, 28, 28), eps_start), cur_anchor, tol=0.001, max_steps=25)
        robust_ub = robust_ub[0]
        print('robustness:', robust_ub)
        fig, name = vis(X, cur_anchor, X_pgd, method, robust_ub, tensorboard=True, plot_curve=True) 
        writer.add_scalar(name+'_robust', robust_ub, plot_step)
        writer.add_figure(name+'_vis', fig, plot_step)
        plot_step += 1


def empty_expl(Xs):
    anchor_maps = np.zeros_like(Xs)
    return anchor_maps


def random_expl(Xs):
    anchor_maps = np.random.rand(*Xs.shape)
    return anchor_maps


def saliency_expl(expl_method, Xs, ys, model):
    """currently only suppor one example at a time"""
    print('Computing explanation by {}...'.format(expl_method))

    assert Xs.shape[0] == 1
    assert ys.shape[0] == 1

    if expl_method == 'SG':
        sg_r = 0.2
        sg_N = 500
        given_expl = 'Grad'
    else:
        sg_r = None
        sg_N = None
        given_expl = None

    anchor_maps, _ = get_explanation_pdt(Xs, model, ys, expl_method, sg_r=sg_r, sg_N=sg_N, given_expl=given_expl, binary_I=False)
    anchor_maps = np.abs(anchor_maps)
    return anchor_maps


def get_explanation_pdt(image, model, label, exp, sg_r=None, sg_N=None, given_expl=None, binary_I=False):
    label = label[0]
    image_v = Variable(image, requires_grad=True)
    model.zero_grad()
    out = model(image_v)
    pdtr = out[:, label]
    pdt = torch.sum(out[:, label])

    if exp == 'Grad':
        pdt.backward()
        grad = image_v.grad
        expl = grad.data.cpu().numpy()
        if binary_I:
            expl = expl * image.cpu().numpy()
    elif exp == 'LOO':        
        mask = torch.eye(784).view(784, 1, 28, 28).bool()
        image_loo = image.repeat(784, 1, 1, 1)
        image_loo[mask] = 0
        image_loo = Variable(image_loo, requires_grad=True)
        pred = model(image_loo)[:, label]
        expl = pred - pdtr.item()
        expl = expl.data.cpu().numpy()
    elif exp == 'SHAP':
        expl = shap(image.cpu(), label, pdt, model, 20000)
    elif exp == 'IG':
        for i in range(10):
            image_v = Variable(image * i/10, requires_grad=True)
            model.zero_grad()
            out = model(image_v)
            pdt = torch.sum(out[:, label])
            pdt.backward()
            grad = image_v.grad
            if i == 0:
                expl = grad.data.cpu().numpy() / 10
            else:
                expl += grad.data.cpu().numpy() / 10
        if binary_I:
            expl = expl * image.cpu().numpy()
    else:
        raise NotImplementedError('Explanation method not supported.')

    return expl, pdtr
