get_feature_importance.py

import sys
import os
import argparse
from torch import backends, cuda, no_grad
import numpy as np
import pandas as pd
from my_utilities.config_reader import read_config
from my_utilities.data_provider_NN import get_data
import torch
from sklearn.utils import shuffle
from torch.utils.data import DataLoader, Dataset
from torchmetrics import Accuracy


DATASET_PATH = "/exploiting_model_multiplicity/data/"
BASEDIR_MODELS = "/exploiting_model_multiplicity/models/"
BASEDIR_PREDICTIONS = "/exploiting_model_multiplicity/experiment_results/"
FILE_SUFFIX = ".csv"


class Data(Dataset):
    def __init__(self, X, y):
        self.index = X.index.values
        X = X.values
        y = y.values
        self.X = torch.from_numpy(X.astype(np.float32))
        self.y = torch.from_numpy(y).type(torch.LongTensor)
        self.len = self.X.shape[0]

    def __getitem__(self, index):
        return self.X[index], self.y[index], self.index[index]

    def __len__(self):
        return self.len


# new dataloder
def get_permuted_loader(
    dataset,
    batch_size: int = 1,
    seed_value: int = 42,
) -> DataLoader:
    torch.manual_seed(seed_value)
    if batch_size == 0:
        batch_size = len(dataset)
    return DataLoader(
        dataset=dataset, batch_size=batch_size, shuffle=False, drop_last=True
    )


def _save_results(df, dataset_name, project_name):
    base_path_to_save = os.path.join(BASEDIR_PREDICTIONS, dataset_name, project_name)
    if not os.path.exists(base_path_to_save):
        os.makedirs(base_path_to_save)
    name = "importances"
    df.to_csv(base_path_to_save + f"/{name}.csv", index=True)


def get_downloaded_models(dataset_name: str, project_name: str) -> list:
    # Unnamed: 0
    path_results = os.path.join(
        BASEDIR_MODELS,
        dataset_name,
        project_name,
        "metadata",
        "all_models_filtered" + FILE_SUFFIX,
    )
    models_filtered = pd.read_csv(path_results)
    print("FILTERED MODELS")
    print(len(models_filtered))
    models_names = [i.split(":")[0] for i in models_filtered["Unnamed: 0"].values]
    print("Models selected in function: ", len(models_names))
    models = []
    for root, dirs, files in os.walk(BASEDIR_MODELS + dataset_name):
        for file in files:
            if (
                project_name in root
                and not file.endswith(".csv")
                and file in models_names
            ):
                path = os.path.join(root, file)
                models.append(path)
    return models


def get_baseline_score(dataloader, list_of_models):
    df = {}
    print(f"Selected {len(list_of_models)} models")
    models_index = [path.split("/")[-1] for path in list_of_models]
    device = (
        "cuda"
        if cuda.is_available()
        else "mps"
        if backends.mps.is_available()
        else "cpu"
    )
    num_iters = len(dataloader)
    for model_path, model_name in zip(list_of_models, models_index):
        accuracy = Accuracy(task="binary")
        fin_accuracy = 0.0
        cuda.empty_cache()
        model = torch.jit.load(model_path)
        model.eval()
        for data in dataloader:
            inputs, labels, _ = data
            inputs, labels = inputs.to(device), labels.to(device)
            with no_grad():
                outputs = model(inputs)
                fin_accuracy += accuracy(outputs, labels)
        original_metric = fin_accuracy.item() / num_iters
        print(f"For model {model_name} original metric {original_metric}")
        df[model_name] = original_metric
    return df


def get_features_importance(list_of_models, test_dir, permutations, original_metrics):
    X_val = pd.read_csv(f"{test_dir}/X.csv", index_col=0)
    y_val = pd.read_csv(f"{test_dir}/y.csv", index_col=0)
    models_index = [path.split("/")[-1] for path in list_of_models]
    all_importance = pd.DataFrame(index=X_val.columns)
    num_features = len(X_val.columns) - 1
    num_models = len(list_of_models)
    models_processed = 0
    for model_path, model_name in zip(list_of_models, models_index):
        models_processed += 1
        print(f"Model asessed: {model_name}")
        cuda.empty_cache()
        model = torch.jit.load(model_path)
        model.eval()
        feature_importance = {}
        for feature_idx in range(X_val.shape[1]):
            print(
                f"Models {models_processed}/{num_models}, features {feature_idx}/{num_features}"
            )
            permuted_metric_sum = 0
            feature_name = X_val.iloc[:, feature_idx].name
            # several permutations
            for _ in range(permutations):
                accuracy = Accuracy(task="binary")
                permuted_accuracy = 0.0
                X_permuted = X_val.copy()
                X_permuted.iloc[:, feature_idx] = shuffle(
                    X_permuted.iloc[:, feature_idx].values
                )
                permuted_dataset = Data(X_permuted, y_val)
                permuted_loader = get_permuted_loader(permuted_dataset, batch_size=1)
                num_iters = len(permuted_loader)
                with torch.no_grad():
                    for inputs, labels, _ in permuted_loader:
                        outputs = model(inputs)
                        permuted_accuracy += accuracy(outputs, labels)
                    fin_perm_acc = permuted_accuracy / num_iters

                permuted_metric_sum += fin_perm_acc

            feature_importance[feature_name] = (
                original_metrics[model_name] - permuted_metric_sum.item() / permutations
            )
        all_importance[model_name] = feature_importance.values()
    return all_importance


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Getting predictions for the experiment"
    )
    parser.add_argument(
        "--config", help="Config/Experiment name", type=str, default=None
    )
    parser.add_argument(
        "--permutations", help="Number of permutations", type=int, default=1
    )
    args = parser.parse_args()
    permutations = args.permutations
    config = read_config(args.config)
    dataset_name = config.get("parameters").get("dataset").get("value")
    project_name = config.get("project")
    test_dir = DATASET_PATH + f"{dataset_name}" + "/" + "test/"
    val_loader = get_data(test_dir, batch_size=1)
    print(f"For {dataset_name} of experiment {project_name} predictions started.")

    list_of_models = get_downloaded_models(dataset_name, project_name)
    print("FUNCTION RESULT len: ", len(list_of_models))

    original_metrics = get_baseline_score(val_loader, list_of_models)
    print("original_metrics: ", original_metrics)
    importances = get_features_importance(
        list_of_models, test_dir, permutations, original_metrics
    )

    _save_results(importances, dataset_name, project_name)