analyze_imagenet_hpo.py

from collections import defaultdict
import argparse
import os

from torchvision.transforms import InterpolationMode, Resize
import captum
import torch
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec

from pnpxai.core.modality.modality import ImageModality
from pnpxai.explainers import LRPUniformEpsilon, IntegratedGradients, KernelShap
from pnpxai.explainers.utils.feature_masks import FeatureMaskFunction
from pnpxai.explainers.utils.function_selectors import FunctionSelector
from pnpxai.explainers.utils.postprocess import PostProcessor
from pnpxai.evaluator.metrics import MoRF, LeRF, AbPC
from pnpxai.evaluator.optimizer import Objective, optimize

from experiments.utils import (
    set_seed,
    get_torchvision_model,
    get_imagenet_sample_from_hf,
    denormalize_image,
    save_pickle_data,
    load_pickle_data,
)


# plot settings
plt.rcParams['font.family'] = 'Times New Roman'


# configs
TARGET_EXPLAINERS = {
    'lrpe': {
        'pnpxai': LRPUniformEpsilon,
        'captum': captum.attr.LRP,
        'dname': r'LRP-Uniform$\varepsilon$',
    },
    'ig': {
        'pnpxai': IntegratedGradients,
        'captum': captum.attr.IntegratedGradients,
        'dname': 'Integrated Gradients'
    },
    'ks': {
        'pnpxai': KernelShap,
        'captum': captum.attr.KernelShap,
        'dname': 'KernelSHAP',
    }
}

TARGET_METRICS = {
    'morf': {
        'cls': MoRF,
        'dname': r'MoRF$\downarrow$',
    },
    'lerf': {
        'cls': LeRF,
        'dname': r'LeRF$\uparrow$',
    },
    'abpc': {
        'cls': AbPC,
        'dname': r'AbPC$\uparrow$',
    },
}


# a custom feature mask function
class Checkerboard(FeatureMaskFunction):
    def __init__(
        self,
        size=[20, 20],
    ):
        assert len(size) == 2
        self.size = size
        self._n_checkers = size[0] * size[1]

    def __call__(self, inputs: torch.Tensor):
        assert inputs.dim() == 4

        bsz, c, h, w = inputs.size()
        # print(input_size)

        resize = Resize([h, w], interpolation=InterpolationMode.NEAREST)

        patch_masks = []
        for i in range(self._n_checkers):
            mask = np.zeros(self._n_checkers)
            mask[i] = i
            mask = resize(
                torch.Tensor(mask).reshape(-
                    1,self.size[0], self.size[1])).unsqueeze(1)
            patch_masks.append(mask.numpy())
        return torch.from_numpy(sum(patch_masks)).squeeze(1).repeat(
            bsz, 1, 1).long().to(inputs.device)


def analyze(args, fp_data):
    """
    Performs HPO and saves raw data for a single data instance.
    """

    # Setup
    set_seed(args.seed)
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    model, transform = get_torchvision_model('resnet18')
    model = model.to(device)
    model.eval()

    img, label = get_imagenet_sample_from_hf(
        transform,
        hf_repo_id="geonhyeongkim/imagenet-samples-for-pnpxai-experiments",
        indices=[args.data_id],
    )
    imgs = img.unsqueeze(0).to(device)
    targets = torch.tensor(args.data_id).unsqueeze(0).to(device)

    # Get explanations and evaluations
    fm_selector = FunctionSelector(
        {'checkerboard': Checkerboard})
    modality = ImageModality(   # set data modality
        channel_dim=1,
        feature_mask_fn_selector=fm_selector,
    )

    plot_data = {   # container for collecting plot data
        'img': denormalize_image(img, mean=transform.mean, std=transform.std),
        'label': label,
        'heatmaps': defaultdict(dict),
        'values': defaultdict(lambda: defaultdict(dict)),
    }
    for explainer_key in TARGET_EXPLAINERS:
        # --- pnpxai -----------------------------------------------------------

        # Create explainer
        pnp_explainer = TARGET_EXPLAINERS[explainer_key]['pnpxai'](model)
        metric = TARGET_METRICS['abpc']['cls'](
            model=model, explainer=pnp_explainer,
        )   # metric to be used as objective: AbPC

        # Get default postprocessor to initialize optimization: (SumPos, MinMax)
        default_pp = modality.get_default_postprocessors()[0]

        # Optimize explainer
        obj = Objective(
            explainer=pnp_explainer,
            postprocessor=default_pp,
            metric=metric,
            modality=modality,
            inputs=imgs,
            targets=targets,
        )
        study = optimize(
            obj,
            direction='maximize',
            n_trials=args.n_trials,
            sampler='tpe',
            seed=args.seed,
        )
        opt_explainer, opt_pp = study.best_trial.user_attrs.values()

        # Get explanation of the optimized explainer
        if explainer_key == 'ks':
            setattr(opt_explainer, 'n_samples', 300)
        opt_attrs = opt_explainer.attribute(imgs, targets)
        opt_attrs_pp = opt_pp(opt_attrs)
        plot_data['heatmaps']['pnpxai'][explainer_key] = (
            opt_attrs_pp.squeeze().detach().cpu().numpy())

        
        # --- captum -----------------------------------------------------------

        # Create explainer
        captum_kwargs = {}
        if explainer_key == 'ks':
            # Use same feature mask with pnpxai
            captum_kwargs['feature_mask'] = Checkerboard()(imgs)
            captum_kwargs['n_samples'] = 300
        captum_explainer = TARGET_EXPLAINERS[explainer_key]['captum'](model)

        # Get explanation of captum explainer
        captum_attrs = captum_explainer.attribute(
            inputs=imgs, target=targets, **captum_kwargs)
        captum_attrs_pp = default_pp(captum_attrs)
        plot_data['heatmaps']['captum'][explainer_key] = (
            captum_attrs_pp.squeeze().detach().cpu().numpy())


        # --- evaluation -------------------------------------------------------

        for metric_key in TARGET_METRICS:
            # Create metric and evaluate
            pnp_value = TARGET_METRICS[metric_key]['cls'](
                model=model, explainer=pnp_explainer).evaluate(
                inputs=imgs, targets=targets, attributions=opt_attrs_pp).item()
            plot_data['values']['pnpxai'][explainer_key][metric_key] = pnp_value
            captum_value = TARGET_METRICS[metric_key]['cls'](
                model=model, explainer=captum_explainer).evaluate(
                inputs=imgs, targets=targets, attributions=captum_attrs_pp).item()
            plot_data['values']['captum'][explainer_key][metric_key] = captum_value

    # Save the data
    os.makedirs(os.path.dirname(fp_data), exist_ok=True)
    save_pickle_data(data=plot_data, filepath=fp_data)


DEFAULT_SELECTED_SAMPLE_INDICES = [
    75, 358, 367, 852,
]


def visualize(args, fp_data, fp_fig):
    """
    Loads saved data for a single instance and generates visualization.
    """
    # Load the data
    plot_data = load_pickle_data(fp_data)

    # Set layout
    edge_size = 2.5
    fig = plt.figure(figsize=(edge_size*3, edge_size*3))
    outer = gridspec.GridSpec(
        2, 1,
        height_ratios=[1.0, 2.4],
    )
    row1 = gridspec.GridSpecFromSubplotSpec(
        1, 3,
        subplot_spec=outer[0],
    )
    ax_img = fig.add_subplot(row1[0, 0])
    ax_bar = fig.add_subplot(row1[0, 1:3])
    row23 = gridspec.GridSpecFromSubplotSpec(
        2, 3,
        subplot_spec=outer[1],
        hspace=0.04,
        wspace=0.04
    )
    axes_heatmaps = [[
        fig.add_subplot(row23[i, j]) for j in range(3)] for i in range(2)]

    # Plot sample img
    ax_img.imshow(plot_data['img'])
    ax_img.set_xticks([]); ax_img.set_yticks([])
    ax_img.set_aspect('equal')
    ax_img.set_title(plot_data['label'].replace('_', ' ').title(), fontsize=15)

    # Plot differences in evaluations between pnpxai and captum
    bar_data = defaultdict(list)
    for explainer_key in TARGET_EXPLAINERS:
        for metric_key in TARGET_METRICS:
            diff = (
                plot_data['values']['pnpxai'][explainer_key][metric_key]
                - plot_data['values']['captum'][explainer_key][metric_key]
            )
            bar_data[explainer_key].append(diff)

    x = np.arange(len(TARGET_METRICS))
    width = 0.25
    for i, (explainer_key, evals) in enumerate(bar_data.items()):
        ax_bar.bar(
            x + i*width, evals, width,
            label=TARGET_EXPLAINERS[explainer_key]['dname'])

    ax_bar.set_ylim(-.2, .8)
    ax_bar.set_title('PnPXAI - Captum', fontsize=15)
    ax_bar.set_xticks(
        x + width,
        [TARGET_METRICS[nm]['dname'] for nm in TARGET_METRICS],
        fontsize=10,
    )
    ax_bar.grid(axis='y')
    ax_bar.margins(x=0.01)

    # Plot heatmaps
    for c, explainer_key in enumerate(TARGET_EXPLAINERS):
        alpha = 1.
        if explainer_key == 'ks':
            axes_heatmaps[0][c].imshow(plot_data['img'])
            axes_heatmaps[1][c].imshow(plot_data['img'])
            alpha *= .75

        axes_heatmaps[0][c].imshow(
            plot_data['heatmaps']['captum'][explainer_key],
            cmap='Reds', alpha=alpha,
        )
        axes_heatmaps[1][c].imshow(
            plot_data['heatmaps']['pnpxai'][explainer_key],
            cmap='Reds', alpha=alpha,
        )

        axes_heatmaps[0][c].set_title(
            TARGET_EXPLAINERS[explainer_key]['dname'], fontsize=15)
        if c == 0:
            axes_heatmaps[0][c].set_ylabel('Captum', fontsize=15)
            axes_heatmaps[1][c].set_ylabel(
                'PnPXAI (Ours)', fontsize=15, fontweight='bold')

        axes_heatmaps[0][c].set_xticks([]); axes_heatmaps[0][c].set_yticks([])
        axes_heatmaps[1][c].set_xticks([]); axes_heatmaps[1][c].set_yticks([])
        axes_heatmaps[0][c].set_aspect('equal')
        axes_heatmaps[1][c].set_aspect('equal')

    # Save figure
    fig.legend(
        loc='upper center', ncols=3,
        bbox_to_anchor=(.5, 1.0), frameon=False, fontsize=12)
    os.makedirs(os.path.dirname(fp_fig), exist_ok=True)
    fig.savefig(fp_fig, bbox_inches='tight', pad_inches=0.02, dpi=300)
    print(f"Visualization saved for Data ID {args.data_id}: {fp_fig}")


def main():
    """Main execution function"""
    # Arguments
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_id', type=int, required=True, help='The specific ID of the data instance to process.')
    parser.add_argument('--save_dir', type=str, default='results/analyze_imagenet_hpo/')
    parser.add_argument('--seed', type=int, default=42)
    parser.add_argument('--n_trials', default=100, type=int)
    parser.add_argument('--analyze', action='store_true')
    parser.add_argument('--visualize', action='store_true')
    args = parser.parse_args()

    # Set result filepaths
    os.makedirs(args.save_dir, exist_ok=True)
    fp_data = os.path.join(args.save_dir, 'raw', f'{args.data_id}.pkl')
    fp_fig = os.path.join(args.save_dir, 'figures', f'{args.data_id}.pdf')

    # Run experiment
    if args.analyze:
        analyze(args, fp_data)
    if args.visualize:
        if not os.path.exists(fp_data):
            raise Exception(
                f'{fp_data} not found. Try again with the following flag: --analyze')
        visualize(args, fp_data, fp_fig)


if __name__ == '__main__':
    main()