OODdetector.py

import torch as th
import logging
import reductions as rd
import tqdm
import os
from reductions import StatefulReductionFactory,_DEFAULT_SPATIAL_REDUCTIONS
from utils.meters import PvalueMatcher,PvalueMatcherFromSamples,AverageMeter,SimpleOnlineMeterFactory,MeterDict
from utils.misc import Recorder
from dataclasses import dataclass
from functools import partial
import layer_selection as ls
from channel_selection import ChannelSelect
from typing import Callable, List, Dict

class FunctionComposition():
    def __init__(self, f1, f2):
        self.f1 = f1
        self.f2 = f2

    def __call__(self, x):
        return self.f2(self.f1(x))

#Legacy
PickleableFunctionComposition = FunctionComposition



def gen_inference_fn(ref_stats_dict, reduction_dict=None):
    reduction_dict = reduction_dict or {}
    def _batch_calc(trace_name, m, inputs):
        if type(inputs) != tuple:
            inputs = (inputs,)
        class_specific_stats = [{} for _ in range(len(ref_stats_dict))]
        for reduction_name, reduction_fn in reduction_dict.items():
            shared_reductions_per_input = []
            for i in inputs:
                if isinstance(reduction_fn, StatefulReductionFactory):
                    i_ = reduction_fn(trace_name, create_new=False).shared_reduction(i)
                else:
                    i_ = reduction_fn(i)
                shared_reductions_per_input.append(i_)
            for c, class_stat_dict in enumerate(ref_stats_dict):
                reduction_stat = class_stat_dict[trace_name[:-8]][reduction_name]
                pval_per_input = []
                for e, per_input_stat in enumerate(reduction_stat):
                    ret_channel_strategy = {}
                    assert isinstance(per_input_stat, BatchStatsCollectorRet)
                    # if per_input_stat.reduction_fn != reduction_dict[reduction_name]:
                    #     assert per_input_stat.reduction_fn.f1 == reduction_dict[reduction_name]
                    reduced = shared_reductions_per_input[e]
                    # if isinstance(reduction_fn, StatefulReductionFactory):
                    #     reduction_fn.set_class_input(c,e,create_new=False)
                    #     reduced = reduction_fn(trace_name,create_new=False).specialized_reduction(reduced)
                    for channle_reduction_name, rec in per_input_stat.channel_reduction_record.items():
                        # if per_input_stat.reduction_fn != reduction_dict[reduction_name]:
                        #     # overwrite function to match old measure file with new format, note that channels
                        #     # reduction can differ between classes
                        #     rec['fn'] = FunctionComposition(f1=per_input_stat.reduction_fn.f2,
                        #                                     f2=per_input_stat.reduction_fn.f1)
                        ret_channel_strategy[channle_reduction_name] = rec['fn'](reduced)
                    #
                    # if per_input_stat.reduction_fn != reduction_dict[reduction_name]:
                    #     # this will make sure the fn overwite will only happen once
                    #     per_input_stat.reduction_fn = per_input_stat.reduction_fn.f1
                    pval_per_input.append(ret_channel_strategy)
                class_specific_stats[c][reduction_name] = pval_per_input
        return class_specific_stats

    return _batch_calc


from _collections import defaultdict


def zero_init_dict():
    def zero():
        return 0

    return defaultdict(zero)


def fisher_reduce_all_layers(ref_stats, filter_layer=None, using_ref_record=False, class_id=None, weights={}):
    # this function summarises all layer pvalues using fisher statistic
    # since we may have multiple channel reduction strategies (e.g. simes, cond-fisher) the strategy dict should have
    # a mapping from reduction output to the actual pvalue (in simes this is just the returned value, for fisher we need
    # to calculate the distribution for each layer statistic)
    sum_pval_per_reduction = {}
    if weights:
        total_layers = zero_init_dict()
        sum_weights_given = zero_init_dict()
        number_of_weighted_layers = zero_init_dict()
        for layer_name, layer_stats_dict in ref_stats.items():
            #### TODO this is a hack to strip original name from the recorder, replace this by using a regex for weights
            if layer_name.endswith('_forward_input_fn'):
                layer_name_ = layer_name[:-len('_forward_input_fn')]
            elif layer_name.endswith('_forward_output_fn'):
                layer_name_ = layer_name[:-len('_forward_output_fn')]
            else:
                layer_name_ = layer_name

            if class_id is not None:
                layer_stats_dict = layer_stats_dict[class_id]
            for spatial_reduction_name, record_per_input in layer_stats_dict.items():
                if filter_layer and filter_layer(layer_name, spatial_reduction_name):
                    continue

                ## different reductions may use different number of layers, get total numbers of layers per reduction
                total_layers[spatial_reduction_name] += 1
                if layer_name_ in weights:
                    sum_weights_given[spatial_reduction_name] += weights[layer_name_]
                    number_of_weighted_layers[spatial_reduction_name] += 1

    for layer_name, layer_stats_dict in ref_stats.items():
        if class_id is not None:
            layer_stats_dict = layer_stats_dict[class_id]
        for spatial_reduction_name, record_per_input in layer_stats_dict.items():
            if filter_layer and filter_layer(layer_name, spatial_reduction_name):
                continue
            if spatial_reduction_name not in sum_pval_per_reduction:
                sum_pval_per_reduction[spatial_reduction_name] = {}
                if using_ref_record:
                    channel_reduction_names = record_per_input[0].channel_reduction_record.keys()
                else:
                    channel_reduction_names = record_per_input[0].keys()

                for channel_reduction_name in channel_reduction_names:
                    sum_pval_per_reduction[spatial_reduction_name][channel_reduction_name] = 0.

            weight = 1.0
            if weights:
                #### TODO this is a hack to strip original name from the recorder
                if layer_name.endswith('_forward_input_fn'):
                    layer_name_ = layer_name[:-len('_forward_input_fn')]
                elif layer_name.endswith('_forward_output_fn'):
                    layer_name_ = layer_name[:-len('_forward_output_fn')]
                else:
                    layer_name_ = layer_name

                assert total_layers[spatial_reduction_name] >= number_of_weighted_layers[spatial_reduction_name]
                if total_layers[spatial_reduction_name] == number_of_weighted_layers[spatial_reduction_name]:
                    weight = weights[layer_name_]  # sanity, make sure this means all layers really exists in weights
                    assert sum_weights_given[spatial_reduction_name] == 1.0
                else:
                    assert sum_weights_given[spatial_reduction_name] <= 1.0
                    denum = total_layers[spatial_reduction_name] - number_of_weighted_layers[spatial_reduction_name]
                    default_weight = (1.0 - sum_weights_given[spatial_reduction_name]) / denum
                    weight = weights.get(layer_name_, default_weight)

            # all layer inputs are reduced together for now
            for record in record_per_input:
                if using_ref_record:
                    assert isinstance(record, BatchStatsCollectorRet)
                    for channel_reduction_name, channel_reduction_record in record.channel_reduction_record.items():
                        record = channel_reduction_record['record']
                        if 'pval_matcher' in channel_reduction_record:
                            # need to get pvalues first
                            pval = channel_reduction_record['pval_matcher'](record)
                        else:
                            pval = record
                        # free memory after extracting stats
                        # del channel_reduction_record['record']
                        # if 'meter' in (channel_reduction_record.keys()):
                        #     del channel_reduction_record['meter']
                        sum_pval_per_reduction[spatial_reduction_name][channel_reduction_name] += -2 * th.log(
                            pval) * weight
                    # del record.meter
                else:
                    # running in test mode
                    for channel_reduction_name, pval in record.items():
                        sum_pval_per_reduction[spatial_reduction_name][channel_reduction_name] += -2 * th.log(
                            pval) * weight

    return sum_pval_per_reduction


def extract_output_distribution_single_class(layer_wise_ref_stats, target_percentiles=th.tensor([0.05,
                                                                                                 0.1, 0.2, 0.3, 0.4,
                                                                                                 0.5, 0.6, 0.7, 0.8,
                                                                                                 0.9,
                                                                                                 # decision for fisher is right sided
                                                                                                 0.945, 0.94625, 0.9475,
                                                                                                 0.94875,
                                                                                                 0.95,
                                                                                                 # target alpha upper 5%
                                                                                                 0.95125, 0.9525,
                                                                                                 0.95375, 0.955,
                                                                                                 # add more abnormal percentiles for fusions
                                                                                                 0.97, 0.98, 0.99,
                                                                                                 0.995, 0.999, 0.9995,
                                                                                                 0.9999]),
                                             right_sided_fisher_pvalue=False, filter_layer=None, weights={}):
    def _prep_pval_matcher(sum_pval_per_reduction):
        fisher_pvals_per_reduction = {}
        for spatial_reduction_name, sum_pval_record in sum_pval_per_reduction.items():
            logging.debug(f'\t{spatial_reduction_name}:')
            fisher_pvals_per_reduction[spatial_reduction_name] = {}
            # different channle reduction strategies will have different pvalues
            for channel_reduction_name, sum_pval in sum_pval_record.items():
                # use right tail pvalue since we don't care about fisher "normal" looking pvalues that are closer to 0
                kwargs = {'target_percentiles': target_percentiles}
                if right_sided_fisher_pvalue:
                    kwargs.update({'left_side': False, 'right_side': True})
                fisher_pvals_per_reduction[spatial_reduction_name][channel_reduction_name] = PvalueMatcherFromSamples(
                    samples=sum_pval, **kwargs)
                logging.debug(f'\t\t{channel_reduction_name}:\t mean:{sum_pval.mean():0.3f}\tstd:{sum_pval.std():0.3f}')
        return fisher_pvals_per_reduction

    # reduce all layers (e.g. fisher)
    pvalue_matcher_per_group = []
    sum_pvalues_per_group = []
    if isinstance(filter_layer, ls.GroupWhiteListInclude):
        for g in range(filter_layer.n_groups):
            filter_layer.set_work_group(g)
            logging.debug(f'processing group {g} pvalues: {filter_layer.get_work_group_members()}')
            sum_pval_per_reduction = fisher_reduce_all_layers(layer_wise_ref_stats, filter_layer, using_ref_record=True,
                                                              weights=weights)
            sum_pvalues_per_group.append(sum_pval_per_reduction)
            pvalue_matcher_per_group.append(_prep_pval_matcher(sum_pval_per_reduction))
        # this allows the next iteration to recover the global fisher pvalue matcher and append it last
        filter_layer.set_global_group()

    sum_pval_per_reduction = fisher_reduce_all_layers(layer_wise_ref_stats, filter_layer, using_ref_record=True,
                                                      weights=weights)
    pvalue_matcher_per_group.append(_prep_pval_matcher(sum_pval_per_reduction))
    sum_pvalues_per_group.append(sum_pval_per_reduction)

    return pvalue_matcher_per_group, sum_pvalues_per_group
    # update replace fisher output with pvalue per reduction


def extract_output_distribution(all_class_ref_stats, target_percentiles=th.tensor([0.05,
                                                                                   0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7,
                                                                                   0.8, 0.9,
                                                                                   # decision for fisher is right sided
                                                                                   0.945, 0.94625, 0.9475, 0.94875,
                                                                                   0.95,  # target alpha upper 5%
                                                                                   0.95125, 0.9525, 0.95375, 0.955,
                                                                                   # add more abnormal percentiles for fusions
                                                                                   0.97, 0.98, 0.99, 0.995, 0.999,
                                                                                   0.9995, 0.9999]),
                                right_sided_fisher_pvalue=False, filter_layer=None, LDA_fisher=False, weights={}):
    per_class_pvalue_matcher_record = []
    fisher_pvals_per_reduction_all_classes = []
    for e, class_stats_per_layer_dict in enumerate(all_class_ref_stats):
        logging.debug(f'Constructing H0 Pvalue matchers for fisher statistic of class {e}/{len(all_class_ref_stats)}')
        fisher_pval_matcher_per_reduction, fisher_stat_per_reduction = extract_output_distribution_single_class(
            class_stats_per_layer_dict,
            target_percentiles=target_percentiles,
            right_sided_fisher_pvalue=right_sided_fisher_pvalue,
            filter_layer=filter_layer, weights=weights)
        per_class_pvalue_matcher_record.append(fisher_pval_matcher_per_reduction)

        for g, fisher_group in enumerate(fisher_stat_per_reduction):
            if len(fisher_pvals_per_reduction_all_classes) < len(fisher_stat_per_reduction):
                fisher_pvals_per_reduction_all_classes.append({})
            for spatial_reduction_name, sum_pval_record in fisher_group.items():
                if spatial_reduction_name not in fisher_pvals_per_reduction_all_classes[g]:
                    fisher_pvals_per_reduction_all_classes[g][spatial_reduction_name] = {}
                for channel_reduction_name, sum_pval in sum_pval_record.items():
                    if channel_reduction_name not in fisher_pvals_per_reduction_all_classes[g][spatial_reduction_name]:
                        fisher_pvals_per_reduction_all_classes[g][spatial_reduction_name][channel_reduction_name] = []
                    fisher_pvals_per_reduction_all_classes[g][spatial_reduction_name][channel_reduction_name].extend(
                        sum_pval)

    joint_fisher_pval_matcher = None
    if LDA_fisher:
        joint_fisher_pval_matcher = []
        for g, fisher_group in enumerate(fisher_pvals_per_reduction_all_classes):
            joint_fisher_pval_matcher.append({})
            for spatial_reduction_name, sum_pval_record in fisher_group.items():
                joint_fisher_pval_matcher[g][spatial_reduction_name] = {}
                for channel_reduction_name, sum_pval in sum_pval_record.items():
                    joint_fisher_pval_matcher[g][spatial_reduction_name][
                        channel_reduction_name] = PvalueMatcherFromSamples(
                        th.cat(sum_pval).unsqueeze(1), target_percentiles=(1 - th.logspace(0, -5, 1000)),
                        left_side=False,
                        right_side=True)

    return per_class_pvalue_matcher_record, joint_fisher_pval_matcher

class CommonStatsRecorder(Recorder):
    def __init__(self, model, monitor_outputs=False, **kwargs):
        if monitor_outputs:
            kwargs['recording_mode'] = [Recorder._RECORD_OUTPUT_MODE[1]]
            kwargs['output_fn'] = kwargs['input_fn']
            kwargs['input_fn'] = None
        else:
            kwargs['recording_mode'] = [Recorder._RECORD_INPUT_MODE[1]]
        super().__init__(model, recursive=True, device_modifier='same', **kwargs)


class OODDetector():
    def __init__(self, model, all_class_ref_stats, include_matcher_fn, right_sided_fisher_pvalue=True, LDA_fisher=False,
                 shared_reductions=_DEFAULT_SPATIAL_REDUCTIONS, fisher_layer_weights={}):
        self.LDA_fisher = LDA_fisher
        self.stats_recorder = CommonStatsRecorder(model, include_matcher_fn=include_matcher_fn,
                                                  input_fn=gen_inference_fn(all_class_ref_stats, shared_reductions))

        for rc in all_class_ref_stats:
            all_keys = list(rc.keys())
            for k in all_keys:
                if k not in self.stats_recorder.tracked_modules.keys():
                    del rc[k]
        # gc.collect()

        self.ref_layers = all_keys
        self.test_layers = list(self.stats_recorder.tracked_modules.keys())
        logging.debug(f'OODdetector - reference laters: {self.ref_layers}')
        logging.debug(f'OODdetector - test laters: {self.test_layers}')
        for tl in self.test_layers:
            if tl not in self.ref_layers:
                logging.warning(f'OODdetector - test later not in reference: {tl}')
        if isinstance(include_matcher_fn, ls.GroupWhiteListInclude):
            self.filter_layer = include_matcher_fn
        else:
            self.filter_layer = lambda layer_name, reduction_name: layer_name not in self.test_layers

        self.output_pval_matcher, self.FLDA_output_pval_matcher = extract_output_distribution(all_class_ref_stats,
                                                                                              right_sided_fisher_pvalue=right_sided_fisher_pvalue,
                                                                                              filter_layer=self.filter_layer,
                                                                                              target_percentiles=th.linspace(
                                                                                                  0, 1, 1000),
                                                                                              LDA_fisher=self.LDA_fisher,
                                                                                              weights=fisher_layer_weights)
        self.num_classes = len(all_class_ref_stats)
        self.fisher_layer_weights = fisher_layer_weights

    def set_tracking(self,on=True):
        self.stats_recorder.master_record_enable=on

    # helper function to convert per class per reduction to per reduction per class dictionary
    def _gen_output_dict(self,per_class_per_reduction_record):
        # prepare a dict with pvalues per reduction per sample per class i.e. {reduction_name : (BxC)}
        reduction_stats_collection = {}
        for reduction_name in per_class_per_reduction_record[0].keys():
            reduction_stats_collection[reduction_name] = []
            for class_stats in per_class_per_reduction_record:
                reduction_stats_collection[reduction_name].append(class_stats[reduction_name].cpu())
            reduction_stats_collection[reduction_name] = th.cat(reduction_stats_collection[reduction_name], -1)

        return reduction_stats_collection

    # this function should return pvalues in the format of (Batch x num_classes)
    # todo merge this with extract_output_distribution fisher compute (iterate over tracked modules
    #  instead of record entries)
    def get_fisher(self):
        per_class_record = []
        # reduce all layers (e.g. fisher)
        for class_id in range(self.num_classes):
            sum_pval_per_reduction = fisher_reduce_all_layers(self.stats_recorder.record, class_id=class_id,
                                                              using_ref_record=False, weights=self.fisher_layer_weights)
            fisher_pvals_per_reduction = self._extract_fisher_pvalues(sum_pval_per_reduction, class_id=class_id)

            per_class_record.append(fisher_pvals_per_reduction)

        return self._gen_output_dict(per_class_record)

    def _extract_fisher_pvalues(self, sum_pval_per_reduction, class_id, group_id=-1):
        # update fisher pvalue per reduction
        fisher_pvals_per_reduction = {}
        for reduction_name, sum_pval_record in sum_pval_per_reduction.items():
            for s, sum_pval in sum_pval_record.items():
                fisher_pvals_per_reduction[f'{reduction_name}_{s}'] = \
                    self.output_pval_matcher[class_id][group_id][reduction_name][s](sum_pval)
                if self.LDA_fisher:
                    fisher_pvals_per_reduction[f'{reduction_name}_{s}_LDA_fisher'] = \
                        self.FLDA_output_pval_matcher[group_id][reduction_name][s](sum_pval)

        return fisher_pvals_per_reduction

    def get_fisher_groups(self, combine_fn=rd.calc_simes, groups_filter: ls.GroupWhiteListInclude = None):
        if groups_filter is None and not isinstance(self.filter_layer, ls.GroupWhiteListInclude):
            return self.get_fisher()
        groups_filter = groups_filter or self.filter_layer
        per_class_record = []
        # reduce all layers (e.g. fisher)
        for class_id in range(self.num_classes):
            fisher_pvals_per_reduction = {}
            for g in range(groups_filter.n_groups):
                groups_filter.set_work_group(g)
                logging.debug(f'processing group {g} pvalues: {groups_filter.get_work_group_members()}')
                sum_pval_per_reduction = fisher_reduce_all_layers(self.stats_recorder.record,
                                                                  filter_layer=groups_filter,
                                                                  using_ref_record=False, class_id=class_id,
                                                                  weights=self.fisher_layer_weights)
                for k, v in self._extract_fisher_pvalues(sum_pval_per_reduction, group_id=g, class_id=class_id).items():
                    if k not in fisher_pvals_per_reduction:
                        fisher_pvals_per_reduction[k] = []
                    fisher_pvals_per_reduction[k] += [v]

            ## combine groups using fisher (alternativly this can be combined outside)
            for k in fisher_pvals_per_reduction.keys():
                fisher_pvals_per_reduction[k] = th.cat(fisher_pvals_per_reduction[k], -1)
                if combine_fn:
                    fisher_pvals_per_reduction[k] = combine_fn(fisher_pvals_per_reduction[k])

            per_class_record.append(fisher_pvals_per_reduction)

        return self._gen_output_dict(per_class_record)

    def get_simes(self):
        per_class_record = []
        for class_id in range(self.num_classes):
            pval_per_reduction = {}
            for layer_name, layer_stats_dict in self.stats_recorder.record.items():
                # if filter_layer and filter_layer(layer_name):
                #     continue
                if class_id is not None:
                    layer_stats_dict = layer_stats_dict[class_id]
                for spatial_reduction_name, record_per_input in layer_stats_dict.items():
                    # prepare the registry
                    if spatial_reduction_name not in pval_per_reduction:
                        pval_per_reduction[spatial_reduction_name] = {}
                        channel_reduction_names = record_per_input[0].keys()
                        for channel_reduction_name in channel_reduction_names:
                            pval_per_reduction[spatial_reduction_name][channel_reduction_name] = []
                    # all layer inputs are reduced together
                    for record in record_per_input:
                        for channel_reduction_name, pval in record.items():
                            pval_per_reduction[spatial_reduction_name][channel_reduction_name].append(pval)

            # update fisher pvalue per reduction
            pvals_per_reduction = {}
            for reduction_name, sum_pval_record in pval_per_reduction.items():
                for s, pval in sum_pval_record.items():
                    pvals_per_reduction[f'{reduction_name}_{s}'] = rd.calc_simes(th.cat(pval, -1))

            per_class_record.append(pvals_per_reduction)

        return self._gen_output_dict(per_class_record)


## auxilary data containers
@dataclass()
class BatchStatsCollectorRet():
    def __init__(self, reduction_name: str,
                 reduction_fn=lambda x: x,
                 cov: th.Tensor = None,
                 num_observations: int = 0,
                 meter: AverageMeter = None):
        self.reduction_name = reduction_name
        self.reduction_fn = reduction_fn
        ## collected stats
        self.cov = cov
        self.num_observations = num_observations
        # spatial reduction meter
        self.meter = meter
        # record to hold all information on channel reduction methods
        self.channel_reduction_record = {}


import numpy as np


def get_percentiles_targets(focal_points=[0.025, 0.5, 0.975], min_resolution=1 / 1000, extreme_sub_samples=100):
    # return th.arange(min_resolution,1,min_resolution)
    focal_points += [i / 10 for i in range(1, 10)]
    percentiles = []
    for focal in focal_points:
        percentiles += [focal - min_resolution, focal, focal + min_resolution]
    min_p = min(percentiles)
    max_p = max(percentiles)
    percentiles += np.linspace(min_resolution, min_p, extreme_sub_samples).tolist() + \
                   np.linspace(max_p, 1 - min_resolution, extreme_sub_samples).tolist()
    percentiles.sort()
    return th.tensor(percentiles)


def truncate_percentiles_to_valid_resolution(percentiles, min_resolution):
    return (((percentiles + (min_resolution / 2)) // min_resolution) * min_resolution).unique()


@dataclass()
class BatchStatsCollectorCfg():
    LDA_tracker: Dict = None
    cov_off: bool = True
    # using partial stats for mahalanobis covariance estimate
    partial_stats: bool = True  # False
    update_tracker: bool = True
    find_simes: bool = False
    find_cond_fisher: bool = True
    fisher_cond_thresh: float = 1.0
    mahalanobis: bool = False
    sum_channel_diff: bool = False
    target_percentiles = None
    # target_percentiles = th.tensor([0.001, 0.002, 0.005, 0.01,
    #                                 # estimate more percentiles next to the target alpha
    #                                 0.02, 0.023, 0.024, 0.025, 0.026, 0.027, 0.03,
    #                                 # collect intervals for better layer reduction statistic approximation
    #                                 0.045, 0.047, 0.049, 0.05, 0.051, 0.053, 0.055, 0.07, 0.1, 0.2, 0.3, 0.4,
    #                                 0.5])  # percentiles will be mirrored
    num_edge_samples: int = 200

    def __init__(self, batch_size, reduction_dictionary=None, include_matcher_fn=None,
                 sampled_channels: Dict[str, th.Tensor] = None, target_percentiles=None):
        self.sample_channels = sampled_channels
        # which reductions to use ?
        self.reduction_dictionary = reduction_dictionary or _DEFAULT_SPATIAL_REDUCTIONS
        # which layers to collect?
        self.include_matcher_fn = include_matcher_fn or ls._default_matcher_fn
        if target_percentiles is None:
            self.target_percentiles = get_percentiles_targets(min_resolution=1 / batch_size)
        else:
            self.target_percentiles = target_percentiles
            # self.target_percentiles = th.cat([self.target_percentiles, (1 - self.target_percentiles).sort()[0]])
            ## adjust percentiles to the specified batch size
        self.target_percentiles = truncate_percentiles_to_valid_resolution(self.target_percentiles, 1 / batch_size)

        # assert 0.5 in self.target_percentiles, 'tensor must include median'
        logging.debug(f'measure target percentiles {self.target_percentiles.numpy()}')


# simple loop over measure data to collect statistics
def _loop_over_data(model,loader,device,epochs,min_batch_lim=100):
    model.eval()
    with th.no_grad():
        for _ in tqdm.trange(epochs):
            for d, l in loader:
                if min_batch_lim and d.shape[0] < min_batch_lim:
                    break
                _ = model(d.to(device))


# here we only collect the eCDF for spatial reductions
def measure_data_statistics_part1(loader, model, measure_settings: BatchStatsCollectorCfg,
                                  epochs=5, model_device='cuda', collector_device='same',
                                  ):
    compute_cov_on_partial_stats = measure_settings.partial_stats and not measure_settings.cov_off
    ## bypass the simple recorder dictionary with a meter dictionary to track per layer statistics
    tracker = MeterDict(meter_factory=SimpleOnlineMeterFactory(batched=True, track_percentiles=True, per_channel=True,
                                                               target_percentiles=measure_settings.target_percentiles,
                                                               number_edge_samples=measure_settings.num_edge_samples,
                                                               track_cov=compute_cov_on_partial_stats))

    # function collects statistics of a batched tensors, return the collected statistics per input tensor
    def _batch_stats_collector_part1(trace_name, m, inputs):
        if type(inputs) != tuple:
            inputs = (inputs,)

        for e, i in enumerate(inputs):
            for reduction_name, reduction_fn in measure_settings.reduction_dictionary.items():
                if isinstance(reduction_fn, StatefulReductionFactory):
                    reduction_fn.set_input(e)
                    i_ = reduction_fn(trace_name)(i, measuring=True)
                else:
                    i_ = reduction_fn(i)

                if collector_device != 'same' and collector_device != model_device:
                    i_ = i_.to(collector_device)
                num_observations, channels = i_.shape
                tracker_name = f'{trace_name}_{reduction_name}:{e}'
                tracker.update({tracker_name: i_})

    def _dummy_reducer(old, new):
        return old

    model.to(model_device)
    r = CommonStatsRecorder(model, include_matcher_fn=measure_settings.include_matcher_fn,
                            input_fn=_batch_stats_collector_part1, activation_reducer_fn=_dummy_reducer)

    logging.info(f'\t\tmeasuring {"covariance " if compute_cov_on_partial_stats else ""} mean and percentiles '
                 f'for all spatial reductions')
    _loop_over_data(model,loader,model_device,epochs)
    r.record.clear()
    r.remove_model_hooks()
    return tracker


# here we would like to use previously gathered statistics over spatial reductions to collect layer statistics
# (used to extract layer pvalue by collecting the layer statistic distribution)
def measure_data_statistics_part2(tracker, loader, model, measure_settings: BatchStatsCollectorCfg,
                                  epochs=5, model_device='cuda', collector_device='same'):

    #measure_settings = measure_settings or BatchStatsCollectorCfg(batch_size)

    # function collects statistics of a batched tensors, return the collected statistics per input tensor
    def _batch_stats_collector_part2(trace_name, m, inputs):
        stats_per_input = []
        if type(inputs) != tuple:
            inputs = (inputs,)
        for e, i in enumerate(inputs):
            reduction_specific_record = []
            for reduction_name, reduction_fn in measure_settings.reduction_dictionary.items():
                tracker_name = f'{trace_name}_{reduction_name}:{e}'
                specialized_reduction = None
                if isinstance(reduction_fn, StatefulReductionFactory):
                    reduction_fn.set_input(e)
                    i_ = reduction_fn(trace_name, create_new=False)(i)
                    specialized_reduction = reduction_fn(trace_name, create_new=False).specialized_reduction
                    reduction_fn = reduction_fn(trace_name, create_new=False).shared_reduction
                else:
                    i_ = reduction_fn(i)

                if measure_settings.sample_channels and tracker_name in measure_settings.sample_channels:
                    sample_channels = th.tensor(measure_settings.sample_channels[tracker_name])
                    # we update reduction_fn and leverage BatchStatsCollectorRet keep layer specific modifications
                    # Note that sampling before reduction is more efficient, however we reverse the order to simplify
                    # the case where spatial reductions may change the number of channels
                    # reduction_fn = PickleableFunctionComposition(f1=reduction_fn,f2=ChannelSelect(sample_channels.clone()))
                    sample_channels_fn = ChannelSelect(sample_channels)
                    i_ = sample_channels_fn(i_)
                    if specialized_reduction is None:
                        specialized_reduction = sample_channels_fn
                    else:
                        specialized_reduction = FunctionComposition(f1=specialized_reduction, f2=sample_channels_fn)

                else:
                    sample_channels = None

                if collector_device != 'same' and collector_device != model_device:
                    i_ = i_.to(collector_device)

                num_observations, channels = i_.shape
                reduction_ret_obj = BatchStatsCollectorRet(reduction_name, reduction_fn,
                                                           num_observations=num_observations)

                # save a reference to the meter for convenience
                try:
                    reduction_ret_obj.meter = tracker[tracker_name]
                except:
                    if tracker_name.startswith('model.'):
                        # tracker_name = tracker_name[6:]
                        assert 0, 'please force recompute refs, last layer is probably missing'
                    else:
                        tracker_name = 'model.' + tracker_name
                    reduction_ret_obj.meter = tracker[tracker_name]
                ## typically second phase measurements
                # this requires first collecting reduction statistics (covariance), then in a second pass we can collect
                if measure_settings.mahalanobis:
                    if measure_settings.LDA_tracker and tracker_name in measure_settings.LDA_tracker:
                        cov_tracker = measure_settings.LDA_tracker
                    else:
                        cov_tracker = tracker

                    if sample_channels is not None:
                        mean, inv_cov = tracker[tracker_name].mean[sample_channels], cov_tracker[tracker_name].inv_cov(
                            sample_channels)
                    else:
                        mean, inv_cov = tracker[tracker_name].mean, cov_tracker[tracker_name].inv_cov()

                    mahalanobis_fn = rd.MahalanobisDistance(mean, inv_cov)
                    # reduce all per channels stats to a single score
                    i_m = mahalanobis_fn(i_)
                    # measure the distribution per layer
                    tracker.update({f'{tracker_name}-@mahalabobis': i_m})
                    if specialized_reduction is not None:
                        # update function channel selection for inference time # todo move all fns outside of the loop
                        mahalanobis_fn = FunctionComposition(f1=specialized_reduction, f2=mahalanobis_fn)
                    reduction_ret_obj.channel_reduction_record.update({'mahalanobis':
                                                                       # used for layer fusion (concatinate over all batches)
                                                                           {'record': i_m,
                                                                            ## used to extract the pval from the output of the spatial reduction output
                                                                            # channel reduction transformation
                                                                            'right_side_pval': True,
                                                                            'fn': mahalanobis_fn,
                                                                            # meter for the channel reduction (used to create pval matcher)
                                                                            'meter': tracker[
                                                                                f'{tracker_name}-@mahalabobis'],
                                                                            }
                                                                       })

                if measure_settings.sum_channel_diff:
                    if sample_channels is not None:
                        mean = tracker[tracker_name].mean[sample_channels]
                    else:
                        mean = tracker[tracker_name].mean

                    fn = rd.SumL1ChannelsDiff(mean)
                    # reduce all per channels stats to a single score
                    i_m = fn(i_)
                    # measure the distribution per layer
                    tracker.update({f'{tracker_name}-@l1_sum_diffs': i_m})
                    if specialized_reduction is not None:
                        # update function channel selection for inference time # todo move all fns outside of the loop
                        fn = FunctionComposition(f1=specialized_reduction, f2=fn)
                    reduction_ret_obj.channel_reduction_record.update({'l1_sum_diffs':
                                                                       # used for layer fusion (concatinate over all batches)
                                                                           {'record': i_m,
                                                                            ## used to extract the pval from the output of the spatial reduction output
                                                                            # channel reduction transformation
                                                                            'right_side_pval': True,
                                                                            'fn': fn,
                                                                            # meter for the channel reduction (used to create pval matcher)
                                                                            'meter': tracker[
                                                                                f'{tracker_name}-@l1_sum_diffs'],
                                                                            }
                                                                       })

                if measure_settings.find_simes or measure_settings.find_cond_fisher:
                    if not hasattr(reduction_ret_obj.meter, 'pval_matcher'):
                        p, q = reduction_ret_obj.meter.get_distribution_histogram()
                        if sample_channels is not None:
                            # need to slice pvalues to sampled channels
                            q = q[:, sample_channels]
                        left_side = True
                        right_side = True
                        if trace_name.startswith('output') and sample_channels is not None:
                            ## if this layer is class dependent softmax output (MSP) then we can use a left test to improve power
                            # assert q.min() >= 0.0 and q.max() <= 1.0 and len(sample_channels) == 1
                            left_side = True
                            right_side = False
                            # import pdb;pdb.set_trace()
                        reduction_ret_obj.meter.pval_matcher = PvalueMatcher(percentiles=p, quantiles=q,
                                                                             left_side=left_side,
                                                                             right_side=right_side)
                    # here we first seek the pvalue for the observated reduction value
                    pval = reduction_ret_obj.meter.pval_matcher(i_)

                    if measure_settings.find_simes:
                        i_s = rd.calc_simes(pval)
                        # tracker.update({f'{tracker_name}-@simes_c': i_})
                        simes_fn = FunctionComposition(f1=reduction_ret_obj.meter.pval_matcher, f2=rd.calc_simes)
                        if specialized_reduction is not None:
                            simes_fn = FunctionComposition(f1=specialized_reduction, f2=simes_fn)

                        reduction_ret_obj.channel_reduction_record.update({'simes_c':
                            {
                                'right_side_pval': False,
                                'record': i_s,
                                'fn': simes_fn
                            }
                        })

                    if measure_settings.find_cond_fisher:
                        i_f = rd.calc_cond_fisher(pval, measure_settings.fisher_cond_thresh)
                        # result is not normalized as pvalues, we need to measure the distribution
                        # of this value to return to pval terms
                        tracker.update({f'{tracker_name}-@fisher_c': i_f})
                        fisher_fn = FunctionComposition(f1=reduction_ret_obj.meter.pval_matcher,
                                                        f2=partial(rd.calc_cond_fisher,
                                                                   thresh=measure_settings.fisher_cond_thresh))
                        if specialized_reduction is not None:
                            fisher_fn = FunctionComposition(f1=specialized_reduction, f2=fisher_fn)
                        reduction_ret_obj.channel_reduction_record.update({'fisher_c':
                                                                               {'record': i_f,
                                                                                'meter': tracker[
                                                                                    f'{tracker_name}-@fisher_c'],
                                                                                'right_side_pval': True,
                                                                                'fn': fisher_fn
                                                                                }
                                                                           })

                reduction_specific_record.append(reduction_ret_obj)

            stats_per_input.append(reduction_specific_record)

        return stats_per_input

    # this functionality is used to calculate a more accurate covariance estimate
    def _batch_stats_reducer_part2(old_record, new_entry):
        stats_per_input = []
        for input_id, reduction_stats_record_n in enumerate(new_entry):
            reductions_per_input = []
            for reduction_id, new_reduction_ret_obj in enumerate(reduction_stats_record_n):
                reduction_ret_obj = old_record[input_id][reduction_id]
                assert reduction_ret_obj.reduction_name == new_reduction_ret_obj.reduction_name
                # aggregate all observed channel reduction values per method
                for channel_reduction_name in new_reduction_ret_obj.channel_reduction_record.keys():
                    reduction_ret_obj.channel_reduction_record[channel_reduction_name]['record'] = \
                        th.cat([reduction_ret_obj.channel_reduction_record[channel_reduction_name]['record'],
                                new_reduction_ret_obj.channel_reduction_record[channel_reduction_name]['record']])

                reductions_per_input.append(reduction_ret_obj)
            stats_per_input.append(reductions_per_input)
        return stats_per_input

    model.to(model_device)
    r = CommonStatsRecorder(model, include_matcher_fn=measure_settings.include_matcher_fn,
                            input_fn=_batch_stats_collector_part2, activation_reducer_fn=_batch_stats_reducer_part2)

    logging.info(f'\t\tcalculating layer pvalues using measured mean and quantiles')
    _loop_over_data(model,loader,model_device,epochs)

    ## build reference dictionary with per layer information per reduction (reversing collection order)
    ret_stat_dict = {}
    for k in r.tracked_modules.keys():
        ret_stat_dict[k] = {}
        for kk, stats_per_input in r.record.items():
            if kk.startswith(k):
                for inp_id, reduction_records in enumerate(stats_per_input):
                    for reduction_record in reduction_records:
                        assert isinstance(reduction_record,BatchStatsCollectorRet)
                        # #todo create a channel reduction pval matcher right here
                        for channel_reduction_entry in reduction_record.channel_reduction_record.values():
                            channel_reduction_entry['record'] = channel_reduction_entry['record'].cpu()
                            if 'meter' in channel_reduction_entry:
                                p, q = channel_reduction_entry['meter'].get_distribution_histogram()
                                # defalut is to use two sided test
                                left_side = True
                                right_side = True
                                if channel_reduction_entry['right_side_pval']:
                                    # for fisher statistic we can use right sided instead
                                    left_side = False
                                    right_side = True
                                pval_matcher = PvalueMatcher(quantiles=q, percentiles=p,
                                                             right_side=right_side,
                                                             left_side=left_side)
                                channel_reduction_entry['pval_matcher'] = pval_matcher
                                # create the final function to retrive the layer pvalue from a given spatial reduction
                                channel_reduction_entry['fn'] = FunctionComposition(channel_reduction_entry['fn'],
                                                                                    pval_matcher)

                        if reduction_record.reduction_name in ret_stat_dict[k]:
                            ret_stat_dict[k][reduction_record.reduction_name] += [reduction_record]
                        else:
                            ret_stat_dict[k][reduction_record.reduction_name] = [reduction_record]
    r.record.clear()
    r.remove_model_hooks()
    return ret_stat_dict


def measure_v2(model, args, measure_cache_part1=None, measure_dataset_part1=None, measure_dataset_part2=None,
               drop_last=False, num_workers=0):
    spatial_reductions = args.spatial_reductions
    measure_dataset_part2 = measure_dataset_part2 or measure_dataset_part1

    def _rotate_class_forall_stateful_factories(c):
        for reduction_name, fn in spatial_reductions.items():
            if isinstance(fn, StatefulReductionFactory):
                fn.set_class(c)

    def _prime_ds_for_per_class_iteration(measure_ds, part2=False):
        if not hasattr(measure_ds, 'classes'):
            measure_ds.classes = list(range(args.num_classes))

        if args.measure_joint_distribution:
            classes = ['joint_distribution']
        else:
            classes = measure_ds.classes.copy()
        if args.LDA and not part2 and 'joint_distribution' != classes[-1]:
            classes += ['joint_distribution']

        targets = measure_ds.targets if hasattr(measure_ds, 'targets') else measure_ds.labels
        if type(targets)!=th.Tensor:
            targets = th.tensor(targets)

        return measure_ds,classes,targets

    if not args.recompute and measure_cache_part1 is not None and os.path.exists(measure_cache_part1):
        logging.info(f'loading cached class statistics (first measure step) from file: {measure_cache_part1}')
        all_class_stat_trackers = th.load(measure_cache_part1,
                                          map_location=args.device if args.collector_device == 'same' else args.collector_device)
    else:
        logging.info(f'Measure part 1')
        all_class_stat_trackers = []
        measure_ds, classes, targets = _prime_ds_for_per_class_iteration(measure_dataset_part1)
        for class_id, class_name in enumerate(classes):
            logging.info(f'\t{class_id}/{len(classes)}\tcollecting stats for class {class_name}')
            ds_ = measure_ds
            if class_name != 'joint_distribution':
                class_ids, = th.where(targets == class_id)
                ds_ = th.utils.data.Subset(measure_ds,class_ids)

            sampler = None  # th.utils.data.RandomSampler(ds_,replacement=True,num_samples=epochs*args.batch_size)
            adjusted_batch_size_measure = min(args.batch_size_measure, len(ds_))
            train_loader = th.utils.data.DataLoader(
                ds_, sampler=sampler,
                batch_size=adjusted_batch_size_measure, shuffle=False if sampler else True,
                num_workers=num_workers, pin_memory=False, drop_last=drop_last)
            _rotate_class_forall_stateful_factories(class_id)
            measure_settings = BatchStatsCollectorCfg(adjusted_batch_size_measure,
                                                      reduction_dictionary=spatial_reductions,
                                                      # todo: maybe split measure and test include fn in Settings
                                                      include_matcher_fn=args.include_matcher_fn_measure)
            if len(ds_) <= args.batch_size_measure:
                measure_settings.num_edge_samples = 0

            ## disable per-class covariance compute for speedup measuring and reduce memory usage
            if (args.LDA and class_name != 'joint_distribution') or not args.mahalanobis_ch_reduction:
                measure_settings.cov_off = True
            # collect basic reduction stats
            class_stats = measure_data_statistics_part1(train_loader, model, epochs=5 if args.augment_measure else 1,
                                                        model_device=args.device,
                                                        collector_device=args.collector_device,
                                                        measure_settings=measure_settings)
            all_class_stat_trackers.append(class_stats)
        if measure_cache_part1 is not None and not args.recompute:
            assert type(measure_cache_part1) == str
            th.save(all_class_stat_trackers, measure_cache_part1)

    if args.LDA:
        assert len(all_class_stat_trackers) == len(measure_ds.classes) + 1
        LDA_tracker = all_class_stat_trackers[-1]
        all_class_stat_trackers = all_class_stat_trackers[:-1]
    else:
        LDA_tracker = None

    if args.channel_selection_fn:
        sampled_channels_dict = args.channel_selection_fn(all_class_stat_trackers)
    else:
        sampled_channels_dict = None

    logging.info(f'Measure part 2')
    all_class_ref_stats = []
    measure_ds_p2, classes, targets = _prime_ds_for_per_class_iteration(measure_dataset_part2, part2=True)

    for class_id, class_name in enumerate(classes):
        logging.info(f'\t{class_id}/{len(classes)}\tcollecting stats for class {class_name}')
        ds_ = measure_ds_p2
        if not args.measure_joint_distribution and class_name != 'joint_distribution':
            class_ids, = th.where(targets == class_id)
            ds_ = th.utils.data.Subset(measure_ds_p2, class_ids)
        adjusted_batch_size_measure = min(args.batch_size_measure, len(ds_))
        sampler = None  # th.utils.data.RandomSampler(ds_,replacement=True,num_samples=epochs*args.batch_size)
        train_loader = th.utils.data.DataLoader(
            ds_, sampler=sampler,
            batch_size=adjusted_batch_size_measure, shuffle=False if sampler else True,
            num_workers=num_workers, pin_memory=False, drop_last=drop_last)
        _rotate_class_forall_stateful_factories(class_id)
        measure_settings = BatchStatsCollectorCfg(adjusted_batch_size_measure, reduction_dictionary=spatial_reductions,
                                                  # todo: maybe split measure and test include fn in Settings
                                                  include_matcher_fn=args.include_matcher_fn_measure,
                                                  sampled_channels=sampled_channels_dict[class_id] if \
                                                      type(sampled_channels_dict) == list else sampled_channels_dict)
        # modify default args for phase2
        measure_settings.LDA_tracker = LDA_tracker
        measure_settings.find_simes = args.simes_ch_reduction
        measure_settings.find_cond_fisher = args.fisher_ch_reduction
        measure_settings.mahalanobis = args.mahalanobis_ch_reduction
        measure_settings.sum_channel_diff = args.sum_ch_diff_reduction

        if len(ds_) <= args.batch_size_measure:
            measure_settings.num_edge_samples = 0

        # collect basic reduction stats
        class_stats = measure_data_statistics_part2(all_class_stat_trackers[class_id], train_loader, model,
                                                    epochs=5 if args.augment_measure else 1,
                                                    model_device=args.device,
                                                    collector_device=args.collector_device,
                                                    measure_settings=measure_settings)
        all_class_ref_stats.append(class_stats)
    return all_class_ref_stats