layer_selection.py

import matplotlib.pyplot as plt
import matplotlib
import seaborn as sns
import scipy.cluster.hierarchy as spc
import numpy as np
import torch as th
import re
from functools import partial

def _default_matcher_fn(n: str, m: th.nn.Module) -> bool:
    return isinstance(m, th.nn.BatchNorm2d) or isinstance(m, th.nn.AvgPool2d) or \
           isinstance(m, th.nn.AdaptiveAvgPool2d) or isinstance(m, th.nn.Identity)  # or isinstance(m, th.nn.Linear)


## default densenet_collector: convolution outputs + avg pool & FC inputs
def include_densenet_layers_fn(n, m):
    return isinstance(m, th.nn.Identity) or isinstance(m, th.nn.AvgPool2d)  #or isinstance(m, th.nn.Linear)


def positional_filter(layer_list, normalized_positions=[1.0]):
    ln = len(layer_list)
    ids = np.unique(np.ceil(ln * np.array(normalized_positions)))
    return [layer_list[int(i)] for i in ids]


def positional_log_filter(layer_list, exp=3):
    ln = len(layer_list)
    sample = np.unique(np.ceil((ln / np.logspace(0, exp, ln))))
    ids = np.min(sample) + ln - sample - 1
    return [layer_list[int(i)] for i in ids]

def slice_layers_filter(layer_list, start, stop, stride=1, keep_last_n=0, tail_first=True):
    ln = len(layer_list)
    if 0 < start < 1:
        start = int(np.floor(ln * start))
    if 0 < stop < 1:
        stop = int(np.ceil(ln * start))
    if 0 < stride < 1:
        stride = int(np.floor(ln * stride))

    if tail_first:
        filtered = layer_list[stop - ln:start - ln:-stride]
    else:
        filtered = layer_list[start:stop:stride]

    if keep_last_n > 0:
        for i in range(ln - keep_last_n, ln):
            if layer_list[i] not in filtered:
                filtered += [layer_list[i]]
    return filtered

# helper functions for more expressive layer filtering
class WhiteListInclude():
    def __init__(self, layer_white_list):
        self.layer_white_list = layer_white_list

    def __call__(self, n, m=None):
        return n in self.layer_white_list


class GroupWhiteListInclude():
    def __init__(self, group_layer_white_reg):
        self.group_layer_white_reg = group_layer_white_reg
        self.all_layers = []
        if isinstance(group_layer_white_reg, dict):
            self.per_reduction_groups = True
            self.n_groups = 0
            for reduction_name, reduction_groups in group_layer_white_reg.items():
                self.n_groups = max(len(reduction_groups), self.n_groups)
                for ll in reduction_groups:
                    self.all_layers += ll
        else:
            self.per_reduction_groups = False
            for ll in group_layer_white_reg:
                self.all_layers += ll
            self.n_groups = len(group_layer_white_reg)
        self.all_layers = list(set(self.all_layers))
        self.set_work_group(0)

    def set_work_group(self, id: int):
        if self.per_reduction_groups:
            self.layer_white_list = {}
            for reduction_name, reduction_groups in self.group_layer_white_reg.items():
                self.layer_white_list[reduction_name] = reduction_groups[id]
        else:
            self.layer_white_list = self.group_layer_white_reg[id]

    def get_work_group_members(self):
        return self.layer_white_list

    def set_global_group(self):
        self.layer_white_list = self.all_layers

    def _match_trace_name(self, n, layer_list):
        if any(re.match(f'^{i}.*', n) for i in layer_list):
            layer_list.append(n)
            return True
        return False

    def __call__(self, n, m):
        if isinstance(m, th.nn.Module):
            return n in self.all_layers

        if self.per_reduction_groups:
            if m not in self.layer_white_list:
                return not (n in self.all_layers or self._match_trace_name(n, self.all_layers))
            return not (n in self.layer_white_list[m] or self._match_trace_name(n, self.layer_white_list[m]))

        return not (n in self.layer_white_list or self._match_trace_name(n, self.layer_white_list))


class LayerSlice(WhiteListInclude):
    def __init__(self, model, include_fn, filter_fn=positional_log_filter):
        all_layers = []
        for n, m in model.named_modules():
            if include_fn(n, m):
                all_layers.append(n)
        super().__init__(layer_white_list=filter_fn(all_layers))


class RegxInclude():
    def __init__(self, pattern):
        self.pattern = pattern

    def __call__(self, n, m):
        # n is the trace name, m is the actual module which can be used to target modules with specific attributes
        return bool(re.fullmatch(self.pattern, n))


# reminder we look at the input of layers, following layers used by mhalanobis paper
densenet_mahalanobis_matcher_fn = WhiteListInclude(
    ['block1', 'block2', 'block3', 'avg_pool', 'model.block1', 'model.block2', 'model.block3', 'model.avg_pool',
     'output'])
resnet_mahalanobis_matcher_fn = WhiteListInclude(
    ['layer1', 'layer2', 'layer3', 'layer4', 'avg_pool', 'model.layer1', 'model.layer2', 'model.layer3', 'model.layer4',
     'model.avg_pool', 'output'])

output_only_fn = WhiteListInclude(['output_softmax', 'output_fc'])


def avg_pool_only(n, m):
    return isinstance(m, th.nn.AvgPool2d)


### 'maxclust' is used to choose number of clusters, 'distance' to choose according to threshold
def findCluster(h0_data, spatial_reduction_name, name_data_set, t=0.8, criterion='distance', plot_layer=False,
                plot_summary=False, channle_reduction_method='simes_c'):

    corr_list = list()
    all_layers = [str(i) for i in h0_data[0].keys()]
    for class_id in range(0, len(h0_data)):
        dim_num = np.array([i for i in range(1, len(all_layers) + 1)])
        full_class = []
        for layer_name in all_layers:
            layer_pval = \
                h0_data[class_id][layer_name][spatial_reduction_name][0].channel_reduction_record[
                    channle_reduction_method][
                'record']
            if 'pval_matcher' in h0_data[class_id][layer_name][spatial_reduction_name][0].channel_reduction_record[
                channle_reduction_method]:
                layer_pval = h0_data[class_id][layer_name][spatial_reduction_name][0].channel_reduction_record[channle_reduction_method]['pval_matcher'](layer_pval)

            if (layer_pval<0).sum():
                print('negative pvalue at',class_id,layer_name,spatial_reduction_name,channle_reduction_method)
            full_class.append(layer_pval)
        full_dat_log = th.log(th.stack(full_class, 1).squeeze(-1)).cpu().numpy()
        corr = np.corrcoef(full_dat_log.T) ## correlation
        if np.isnan(corr).sum() > 0:
            import pdb; pdb.set_trace()
        corr_list.append(corr)
        if plot_layer:
            # fig, axes = plt.subplots(ncols = 2, nrows = 1, sharex = True, figsize = (14, 8), sharey = False)
            #First create the clustermap figure
            clustermap = sns.clustermap(corr, col_cluster=False, linewidth = 0.0, figsize = (12, 8), method = 'complete',
                                        cbar_pos = (1, .2, .03, .4))
            clustermap.fig.suptitle(f'Correlation_{name_data_set}_{spatial_reduction_name}_class_{class_id}')
            # set the gridspec to only cover half of the figure
            clustermap.gs.update(left=0.05, right=0.45)
            #create new gridspec for the right part
            gs2 = matplotlib.gridspec.GridSpec(1, 1, left = 0.6, top = 0.9)
            # create axes within this new gridspec
            ax2 = clustermap.fig.add_subplot(gs2[0])
            # plot boxplot in the new axes
            #axes[0].title.set_text('Correlation between layers')
            for l in range(1, 5, 2):
                fisher_statistic = np.apply_along_axis(lambda x: -sum(x), 1, full_dat_log[:, range(0, full_dat_log.shape[1], l)])
                np.var(fisher_statistic)
                sns.kdeplot(fisher_statistic, shade = True, label = f'each {l} column - var {np.var(fisher_statistic):.3f}', ax = ax2)
            plt.legend()
            plt.show()
            clustermap.savefig(f'Images/Correlation_{name_data_set}_{spatial_reduction_name}_class_{class_id}.png')
    #### Select dimensions according to correlation
    ### Heirarchal clustering
    avg_corr = sum(corr_list) / max(len(corr_list),1)
    pwdist = 1 - abs(avg_corr)
    # take upper half of the distance metric 1-corr
    #pwdist  = spc.distance.pdist(1 - abs(avg_corr)) ### abs for sake of correctness
    pwdist = pwdist[np.triu_indices_from(pwdist, 1)]
    linkage = spc.linkage(pwdist, method='ward')
    # apply thershold to trim connections between weakly correlated layers
    cluster = spc.fcluster(linkage, t = t, criterion = criterion)
    ### Sample from clusters
    chosen_layers = []
    for j in range(1, max(cluster) + 1):
        temp_dims = dim_num[np.where(cluster == j)[0]] - 1
        # take layer with maximum correlation with other layers (avg)
        chosen_dim = temp_dims[np.argmax(avg_corr[temp_dims, :][:, temp_dims].sum(1))]
        chosen_layers.append(chosen_dim)

    if plot_summary:
        fig, axes = plt.subplots(ncols = 2, nrows = 1, figsize = (12, 4), sharey = False, sharex = False)
        fig.suptitle(f'cutoff = {t}, selecting from clusters' , fontsize=14)
        sns.heatmap(avg_corr[chosen_layers, :][:, chosen_layers], ax = axes[0])
        all_fisher_statistic = np.apply_along_axis(lambda x: -sum(x), 1, full_dat_log)
        sns.kdeplot(all_fisher_statistic, label = f'All - var {np.var(all_fisher_statistic):.1f}', shade = True,  ax = axes[1])
        fisher_statistic = np.apply_along_axis(lambda x: -sum(x), 1, full_dat_log[:, chosen_layers])
        sns.kdeplot(fisher_statistic, label = f'{len(chosen_layers)} clusters - var {np.var(fisher_statistic):.1f}, expected - {(len(chosen_layers) / len(dim_num)) * np.var(all_fisher_statistic):.1f} ', shade = True,  ax = axes[1])
        plt.legend()
    return [all_layers[k] for k in chosen_layers], chosen_layers, corr_list, full_dat_log

def findClusterMain(settings, h0_data,cut_off_thres=None,plot=False):
    # corr distance required to consider correlated layers as separate clusters, higher value will lead to less clusters
    cut_off_thres = cut_off_thres or [i / 20 for i in range(0, 20)]
    net = settings.model
    data_set = settings.dataset
    reudction_list = settings.spatial_reductions.keys()
    res_dict = {}
    if plot:
        #### Create variance as function of t plots
        plt.figure(figsize=(14, 8))

    for j in reudction_list:
        result_t_fisher = []
        result_t_conditional = []
        res_dict[j]=[]
        for t in cut_off_thres:
            layer_name, ind, var_corr, full_dat_log = findCluster(h0_data, spatial_reduction_name=j, name_data_set=data_set, t=t,
                                                                  criterion='distance')
            res_dict[j].append(layer_name)
            if plot:
                #### Condiditional fisher output distribution
                fisher_statistic = np.apply_along_axis(lambda x: -sum(x), 1, full_dat_log[:, ind])
                result_t_fisher.append(np.var(fisher_statistic))
                full_dat_log[full_dat_log > np.log(0.1)] = 0
                fisher_statistic = np.apply_along_axis(lambda x: -sum(x), 1, full_dat_log[:, ind])
                result_t_conditional.append(np.var(fisher_statistic))
        if plot:
            plt.plot(cut_off_thres, result_t_fisher, label=f'reduction - {j}, {data_set}, {net}, Fisher')
            plt.plot(cut_off_thres, result_t_conditional, label=f'reduction - {j}, {data_set}, {net}, Fisher_conditional')
            plt.legend()
        # res_dict[i][j] = result_t
        # res_dict[i][j] = layer_name
        # avg_corr = sum(var_corr) / len(var_corr)
        # var_list = [(i - avg_corr)**2 for i in var_corr]
        # ax = plt.axes()
        # sns.heatmap(sum(var_list) / len(var_list), label = 'Variance of correlation matrix')
        # ax.set_title('Variance of correlation matrix')

    # layer_name, ind, var_corr, fisher = findCluster(h0_data, spatial_reduction_name='spatial-max', name_data_set=data_set, t=0.5,
    #                                                 criterion='distance')
    return res_dict


def find_cluster_groups(h0_data, spatial_reduction_name, t=10, criterion='maxclust',
                        channle_reduction_method='simes_c'):
    import scipy.cluster.hierarchy as spc

    corr_list = list()
    all_layers = [str(i) for i in h0_data[0].keys()]
    for class_id in range(0, len(h0_data)):
        dim_num = np.array([i for i in range(len(all_layers))])
        full_class = []
        for layer_name in all_layers:
            layer_pval = \
                h0_data[class_id][layer_name][spatial_reduction_name][0].channel_reduction_record[
                    channle_reduction_method][
                    'record']
            if 'pval_matcher' in h0_data[class_id][layer_name][spatial_reduction_name][0].channel_reduction_record[
                channle_reduction_method]:
                layer_pval = h0_data[class_id][layer_name][spatial_reduction_name][0].channel_reduction_record[
                    channle_reduction_method]['pval_matcher'](layer_pval)

            if (layer_pval < 0).sum():
                print('negative pvalue at', class_id, layer_name, spatial_reduction_name, channle_reduction_method)
            full_class.append(layer_pval)
        full_dat_log = th.log(th.stack(full_class, 1).squeeze(-1)).cpu().numpy()
        corr = np.corrcoef(full_dat_log.T)  ## correlation
        if np.isnan(corr).sum() > 0:
            import pdb;
            pdb.set_trace()
        corr_list.append(corr)
    #### Select dimensions according to correlation
    ### Heirarchal clustering
    avg_corr = sum(corr_list) / max(len(corr_list), 1)
    pwdist = abs(avg_corr)
    # take upper half of the distance metric 1-corr
    # pwdist  = spc.distance.pdist(1 - abs(avg_corr)) ### abs for sake of correctness
    pwdist = pwdist[np.triu_indices_from(pwdist, 1)]
    linkage = spc.linkage(pwdist, method='ward')
    # apply thershold to trim connections between weakly correlated layers
    cluster = spc.fcluster(linkage, t=t, criterion=criterion) - 1
    ### Sample from clusters
    n_clusters = len(np.unique(cluster))
    chosen_layers = [None] * n_clusters
    for j in range(n_clusters):
        chosen_layers[j] = [all_layers[i] for i in (dim_num[np.where(cluster == j)[0]])]
        # take layer with maximum correlation with other layers (avg)
    return chosen_layers

def get_select_layers_from_preset(args,model,ref_stats):
    import logging
    if args.select_layer_mode == 'auto':
        logging.info(f'layer clustering')
        selected_layers_names = findClusterMain(args, ref_stats, cut_off_thres=[0.5])
        # currently we can only look at
        selected_layers_names = selected_layers_names['spatial-max'][0]
        include_matcher_fn_test = WhiteListInclude(selected_layers_names)

    elif args.select_layer_mode == 'auto_group':
        logging.info(f'layer clustering with groups')
        selected_layers_names = {}
        for reduction_name in args.spatial_reductions.keys():
            selected_layers_names[reduction_name] = find_cluster_groups(ref_stats, reduction_name,
                                                                        **args.select_layer_kwargs)
        args.include_matcher_fn_test = GroupWhiteListInclude(selected_layers_names)
    elif args.select_layer_mode == 'logspace':
        logging.info(f'logspace layer selection')
        include_matcher_fn_test = LayerSlice(model, include_fn=args.include_matcher_fn_measure,
                                             filter_fn=partial(positional_log_filter,
                                                               **args.select_layer_kwargs))
        selected_layers_names = include_matcher_fn_test.layer_white_list

    elif args.select_layer_mode == 'positional':
        logging.info(f'positional layer selection')
        include_matcher_fn_test = LayerSlice(model, include_fn=args.include_matcher_fn_measure,
                                             filter_fn=partial(positional_filter, **args.select_layer_kwargs))
        selected_layers_names = include_matcher_fn_test.layer_white_list

    elif args.select_layer_mode == 'slice':
        logging.info(f'slicing layer selection')
        include_matcher_fn_test = LayerSlice(model, include_fn=args.include_matcher_fn_measure,
                                             filter_fn=partial(slice_layers_filter, **args.select_layer_kwargs))
        selected_layers_names = include_matcher_fn_test.layer_white_list

    return selected_layers_names, include_matcher_fn_test