mia.py

"""
Vulnerability estimation code.
"""

import itertools
import warnings
import numpy as np
import pandas as pd

from joblib import Parallel, delayed
from sklearn.ensemble import GradientBoostingClassifier

from utils import log_losses


_MSG_ENFORCE_UNIFORM_PRIOR_AND_MICRODATA = (
    "enforce_uniform_prior has no effect when microdata is set to True."
)


def run_threshold_estimator(
    y_train,
    preds_train,
    y_test,
    preds_test,
    method="average_loss_threshold",
    eval_sets=None,
    microdata=False,
    enforce_uniform_prior=True,
):
    """
    Estimate privacy of a given model using a threshold attack.
    """
    losses_train = log_losses(y_train, preds_train)
    losses_test = log_losses(y_test, preds_test)

    if method == "best_loss_threshold":
        threshold = 0
        best_acc = 0
        for c in itertools.chain(losses_train, losses_test):
            if enforce_uniform_prior:
                # Compute balanced accuracy.
                acc = 0.5 * ((losses_train < c).mean() + (losses_test >= c).mean())
            else:
                acc = np.concatenate([losses_train < c, losses_test >= c]).mean()

            if acc > best_acc:
                threshold = c
                best_acc = acc

    elif method == "average_loss_threshold":
        threshold = losses_train.mean()

    else:
        raise NotImplementedError(method)

    # Eval sets.
    if eval_sets is not None:
        eval_y_train, eval_preds_train, eval_y_test, eval_preds_test = eval_sets
        eval_losses_train = log_losses(eval_y_train, eval_preds_train)
        eval_losses_test = log_losses(eval_y_test, eval_preds_test)
    else:
        eval_losses_train = losses_train
        eval_losses_test = losses_test

    in_guesses = eval_losses_train < threshold
    out_guesses = eval_losses_test >= threshold
    if microdata:
        if enforce_uniform_prior:
            warnings.warn(_MSG_ENFORCE_UNIFORM_PRIOR_AND_MICRODATA)
        index=list(y_train.index) + list(y_test.index)
        return pd.Series(np.concatenate([in_guesses, out_guesses]), index=index)

    else:
        if enforce_uniform_prior:
            return 0.5 * (in_guesses.mean() + out_guesses.mean())
        else:
            return np.concatenate([in_guesses, out_guesses]).mean()


def apply_estimator_func(
    models_y_train,
    models_preds_train,
    models_y_test,
    models_preds_test,
    estimator_func,
    parallel=True,
    microdata=False,
    model_label="model_id",
    example_label="example_id",
    verbose=0,
    n_jobs=8,
    **kwargs,
):
    """Independently apply estimator function to multiple model."""
    it = [
        (models_y_train[i], models_preds_train[i], models_y_test[i], models_preds_test[i])
        for i in range(len(models_y_train))
    ]

    if parallel:
        vulns = Parallel(n_jobs=n_jobs, verbose=verbose)(delayed(estimator_func)(
            *params, microdata=microdata, **kwargs) for params in it)
    else:
        vulns = []
        for params in it:
            vulns.append(estimator_func(*params,  microdata=microdata, **kwargs))

    # Build a dataframe if microdata is on.
    if microdata:
        vulns_extended = []
        targets_extended = []
        index_extended = []
        memberships_extended = []
        for i, vuln in zip(range(len(models_y_train)), vulns):
            vuln = pd.Series(vuln)
            vulns_extended.extend(vuln.values)
            index_extended.extend(vuln.index)
            targets_extended.extend([i] * len(vuln))
            memberships_extended.extend(
                [1] * len(models_preds_train[0]) + [0] * len(models_preds_test[0]))

        multi_index = pd.MultiIndex.from_tuples(
            zip(targets_extended, index_extended),
            names=[model_label, example_label]
        )
        return pd.DataFrame(dict(
            vuln=vulns_extended,
            m=memberships_extended
        ), index=multi_index)

    # Otherwise, return an array of vulnerability estimates per model.
    else:
        return np.array(vulns)


def run_shadow_model_attack(
    shadow_models_y_train,
    shadow_models_preds_train,
    shadow_models_y_test,
    shadow_models_preds_test,
    target_models_y_train,
    target_models_preds_train,
    target_models_y_test,
    target_models_preds_test,
    method="shadow_attack_loss",
    parallel=True,
    microdata=False,
    enforce_uniform_prior=True,
):
    """
    Run a shadow model attack.

    Fit one attack model on part of data; evaluate independently on each target model.
    """
    if microdata and enforce_uniform_prior:
        warnings.warn(_MSG_ENFORCE_UNIFORM_PRIOR_AND_MICRODATA)

    preds_in_learn = np.concatenate(shadow_models_preds_train)
    preds_out_learn = np.concatenate(shadow_models_preds_test)

    # Attacker's features are loss values.
    if method == "shadow_attack_loss":
        y_in_learn = np.concatenate(shadow_models_y_train)
        y_out_learn = np.concatenate(shadow_models_y_test)

        losses_in_learn = log_losses(y_in_learn, preds_in_learn)
        losses_out_learn = log_losses(y_out_learn, preds_out_learn)

        X_learn = np.expand_dims(np.concatenate([losses_in_learn, losses_out_learn]), 1)
        y_learn = np.concatenate([[1] * len(losses_in_learn), [0] * len(losses_out_learn)])

    # Attacker's features are model's outputs.
    elif method == "shadow_attack_preds":
        X_learn = np.expand_dims(np.concatenate([preds_in_learn, preds_out_learn]), 1)
        y_learn = np.concatenate([[1] * len(preds_in_learn), [0] * len(preds_out_learn)])

    else:
        raise NotImplementedError

    # Attacker's classifier is GBDT.
    attacker = GradientBoostingClassifier()
    attacker.fit(X_learn, y_learn)

    def eval_one_model(y_train, preds_train, y_test, preds_test, microdata=False):
        preds_in_eval = preds_train
        preds_out_eval = preds_test

        if method == "shadow_attack_loss":
            losses_in_eval = log_losses(y_train, preds_train)
            losses_out_eval = log_losses(y_test, preds_test)

            X_eval_in = np.expand_dims(losses_in_eval, 1)
            X_eval_out = np.expand_dims(losses_out_eval, 1)
            # y_eval = np.concatenate([[1] * len(losses_in_eval), [0] * len(losses_out_eval)])

        elif method == "shadow_attack_preds":
            X_eval_in = np.expand_dims(preds_in_eval, 1)
            X_eval_out = np.expand_dims(preds_out_eval, 1)
            # y_eval = np.concatenate([[1] * len(preds_in_eval), [0] * len(preds_out_eval)])

        if microdata:
            index=list(y_train.index) + list(y_test.index)
            return pd.Series(np.concatenate(
                attacker.predict(X_eval_in) == np.array([1] * len(X_eval_in)),
                attacker.predict(X_eval_out) == np.array([0] * len(X_eval_out))
            ), index=index)
        else:
            if enforce_uniform_prior:
                return 0.5 * (
                    (attacker.predict(X_eval_in) == np.array([1] * len(X_eval_in))).mean() + \
                    (attacker.predict(X_eval_out) == np.array([0] * len(X_eval_out))).mean()
                )
            else:
                return np.concatenate([
                    attacker.predict(X_eval_in) == np.array([1] * len(X_eval_in)),
                    attacker.predict(X_eval_out) == np.array([0] * len(X_eval_out))
                ]).mean()

    return apply_estimator_func(
        target_models_y_train,
        target_models_preds_train,
        target_models_y_test,
        target_models_preds_test,
        eval_one_model,
        parallel=parallel,
        microdata=microdata,
    )