model_extraction.py

from sklearn.preprocessing import RobustScaler
import joblib
import argparse 
import os 
import numpy as np 

from sampling_strategy import entropy_sampling, random_sampling, mc_dropout, k_center, ensemble
from utils import init_scores, log_and_score
from sklearn_extra.cluster import KMedoids
import pandas as pd
from datetime import datetime
from mlflow import mlflow, log_param

from target import LGBTarget, FileBasedTarget, TorchTarget
from attacker import LGBAttacker, KerasAttacker, KerasDualAttacker, SVMAttacker
from datasets import get_ember_data, get_av_data, get_sorel_data

# MLFLOW: Set to False if not supported
MLFLOW = True

if MLFLOW:
    mlflow.set_tracking_uri('http://localhost:5000')
    mlflow.set_experiment("Experiment X")

prog = "Model extraction using active learning techniques"
parser = argparse.ArgumentParser(prog=prog)
parser.add_argument("-d", "--data_dir", type=str, default=".", required=True, help="Directory that holds the data")
parser.add_argument("-s", "--seed", type=int, default="42", required=False, help="Seed for random states")
parser.add_argument("-m", "--method", type=str, required=False, default="entropy",help="entropy, random, medoids, mc_dropout, k-center, ensemble")
parser.add_argument("-n", "--num_queries", type=int, default="10", required=False, help="Number of query rounds")
parser.add_argument("-b", "--budget", type=int, default="25000", required=False, help="Total query budget")
parser.add_argument("-e", "--num_epochs", type=int, default="1", required=False, help="Number of training epochs per round")
parser.add_argument("-t", "--type", type=str, default="LGB", required=False, choices=["DNN", "dualDNN", "LGB", "SVM"], help="Type of surrogate model")
parser.add_argument("-l", "--log_dir", type=str, default="./logs", required=False, help="Where to store the log files with the results")
parser.add_argument("-tg", "--target_model", type=str, default="ember", required=False, 
    choices=["ember", "sorel-FCNN", "sorel-LGB", "AV1", "AV2", "AV3", "AV4"], help="Target model")
parser.add_argument("-f", "--family", type=str, required=False, default="top10families", 
    choices=["top10families", "Adload", "WannaCry", "Pykse", "Azorult", "Bancteian", "Emotet", "Swisyn", "Vobfus"],
    help="Select top10 families or one specific malware family")
parser.add_argument("--dataset", type=str, choices=["ember", "sorel", "AV"], default="ember", required=False, help="Thief and test dataset")
parser.add_argument("--fpr", type=np.float32, default=0.01, required=False, help="FPR level for surrogate merics.")

parser.set_defaults(balanced=False)

args = parser.parse_args()
data_dir = args.data_dir
seed = args.seed
sampling_method = args.method
num_queries = args.num_queries
budget = args.budget
num_epochs = args.num_epochs
model_type = args.type
log_dir = args.log_dir
target = args.target_model
family = args.family
fpr_target = args.fpr
dataset = args.dataset

BUDGET = budget
SEED_SIZE =  BUDGET // 10
VAL_SIZE = BUDGET // 5

num_instances = (BUDGET - VAL_SIZE - SEED_SIZE) // num_queries
print("Num instances per query", num_instances)

if MLFLOW:
    mlflow.start_run()
    log_param("target", target)
    log_param("model_type", model_type)
    log_param("sampling_method", sampling_method)
    log_param("num_queries", num_queries)
    log_param("num_instances", num_instances)
    log_param("seed", seed)
    log_param("num_epochs", num_epochs)
    log_param("budget", BUDGET)
    log_param("family", family)
    log_param("fpr_target", fpr_target)

# Set the numpy random seed
rg = np.random.default_rng(seed)

if model_type in ['DNN', 'SVM', 'dualDNN']:
    scaling_needed = True
else:
    scaling_needed = False

if model_type == "dualDNN":
    dual = True
else:
    dual = False

if "AV" in target:
    isAV = True
else:
    isAV = False


print("Reading vectorized data...")
if isAV:
    X_train, X_test, y_train_target, y_test_target, y_train, y_test = get_av_data(data_dir, family, target, seed=seed)    
else:
    if dataset == "ember":
        X_train, X_test, y_train, y_test = get_ember_data(data_dir, seed=seed)
    else:
        X_train, X_test, y_train, y_test = get_sorel_data(data_dir, seed=seed)

print("Loading target model and retrieving target labels...")
if target == "ember":
    target_model = LGBTarget(os.path.join(data_dir, 'ember_model_2018.txt'), thresh=0.8336, name=target)
elif target == "sorel-LGB":
    sorel_model = f'/data/mari/sorel-models/lightGBM/seed0/lightgbm.model'
    target_model = LGBTarget(sorel_model, thresh=0.5, name=target)
elif target == "sorel-FCNN":
    model_dir = f'/data/mari/sorel-models/FFNN/seed4/'
    checkpoint_file = os.path.join(model_dir, 'epoch_10.pt')
    target_model = TorchTarget(checkpoint_file, 0.5, target)
elif isAV:
    target_model = FileBasedTarget(None, name=target, labels=y_train_target)
else:
    print("This target is not implemented yet.")
    exit()

# Get target model predictions for measuring disagreement
if isAV:
    y_pred_target = y_test_target
else:
    y_pred_target = target_model(X_test)

# Select the seed data from the unlabeled dataset
idx = rg.choice(X_train.shape[0], SEED_SIZE, replace=False)
X_seed = X_train[idx]
X_train = np.delete(X_train, idx, axis=0)

# Add the correct labels
y_seed_true = y_train[idx]
y_train = np.delete(y_train, idx, axis=0)

# Get the seed labels from the Oracle (target model)
if isAV:
    y_seed = target_model(idx)
else:
    y_seed = target_model(X_seed)
y_seed = np.reshape(y_seed, (SEED_SIZE,))

# Select validation set from the unlabeled data
idx = rg.choice(X_train.shape[0], VAL_SIZE, replace=False)
X_val = X_train[idx]
X_train = np.delete(X_train, idx, axis=0)
# print(X_train.shape)

# Add the correct labels
y_val_true = y_train[idx]
y_train = np.delete(y_train, idx, axis=0)

# Get the validation labels from the Oracle (target model)
if isAV:
    y_val = target_model(idx)
else:
    y_val = target_model(X_val)
y_val = np.reshape(y_val, (VAL_SIZE,))

# Scale the input data if substitute is DNN
# if scaling_needed and not os.path.isfile(f'robust_scaler_{seed}.joblib'):
if scaling_needed:
    print("Fitting the scaler...")
    scaler = RobustScaler()
    scaler.fit(X_train)
    joblib.dump(scaler, os.path.join(log_dir, f'robust_scaler_{seed}.joblib'))

    print("Scaling and clipping the seed data...")
    # scaler = joblib.load(f'robust_scaler_{seed}.joblib')
    X_seed = scaler.transform(X_seed)
    X_seed = np.clip(X_seed, -5, 5)

    print("Scaling and clipping unlabeled and test data...")
    X_test = scaler.transform(X_test)
    X_test = np.clip(X_test, -5, 5)
    X_val = scaler.transform(X_val)
    X_val = np.clip(X_val, -5, 5)
    
    X_train_s = scaler.transform(X_train)
    X_train_s = np.clip(X_train_s, -5, 5)

print("shape of X_seed: ", X_seed.shape)
print("shape of y_seed: ", y_seed.shape)
print("Test data shape:", X_test.shape)
print("Unlabelled data shape:", X_train.shape)

# Take the seed data as cluster centers for k-center
if sampling_method == "k-center":
    X_cluster = X_seed

# Create substitute model
if model_type == 'DNN':
    sub_model = KerasAttacker(early_stopping=30, seed=seed)
    sub_model.train_model(X_seed, y_seed, X_val, y_val, num_epochs)
elif model_type == 'dualDNN':
    sub_model = KerasDualAttacker(early_stopping=30, seed=seed)
    sub_model.train_model(X_seed, y_seed, y_seed_true, 
                        X_val, y_val, y_val_true, num_epochs)
elif model_type == "SVM":
    sub_model = SVMAttacker(seed=seed, max_iter=1000)
    sub_model.train_model(X_seed, y_seed)
else:
    sub_model = LGBAttacker(seed=42)
    sub_model.train_model(X_seed, y_seed, X_val, y_val, boosting_rounds=100, early_stopping=15)

if dual:
    y_proba = sub_model(X_test, y_test)   
else: 
    y_proba = sub_model(X_test)
y_pred = np.round(y_proba)
 
scores = init_scores()
scores["nums"].append(SEED_SIZE+VAL_SIZE)
scores = log_and_score(y_proba, y_pred_target, y_test, scores, fpr_target, logging=True)

for idx in range(num_queries): 
    print('Query no. %d' % (idx + 1))
    # Support entropy and random at the moment
    if sampling_method == 'entropy':
        if scaling_needed:
            query_idx = entropy_sampling(sub_model, X_train_s, y_train, n_instances=num_instances, dual=dual)
        else:
            query_idx = entropy_sampling(sub_model, X_train, y_train, n_instances=num_instances)
    elif sampling_method == 'random':
        query_idx = random_sampling(X_train, rg, n_instances=num_instances)
    elif sampling_method == "medoids":
        # First use entropy to select a number of data points 
        # because k-medoids does not scale
        if scaling_needed:
            q_idx = entropy_sampling(sub_model, X_train_s, y_train, n_instances=10000, dual=dual)
            X_med = X_train_s[q_idx]
        else:
            q_idx = entropy_sampling(sub_model, X_train, y_train, n_instances=10000)
            X_med = X_train[q_idx]
        
        kmed = KMedoids(n_clusters=num_instances, init='k-medoids++', random_state=seed)
        # kmed = KMedoids(n_clusters=num_instances, random_state=seed)
        kmed.fit(X_med)
        query_idx = q_idx[kmed.medoid_indices_]

        # Notes: 1. PCA did not help further
        # 2. Using the DNN output as feature vector did not help much either
        # features = tf.keras.Model(inputs=sub_model.input,
        #                               outputs=sub_model.get_layer("final_dense").output)
        
    elif sampling_method == 'k-center':

        # Greedy k-center as in the paper and Google repo 
        # WARNING: Very slow!  
        if scaling_needed:
            q_idx = entropy_sampling(sub_model, X_train_s, y_train, n_instances=10000, dual=dual)
            X_med = X_train_s[q_idx]
        else:
            q_idx = entropy_sampling(sub_model, X_train, y_train, n_instances=10000)
            X_med = X_train[q_idx]

        q_idx2, X_cluster = k_center(X_cluster, X_med, num_instances)
        query_idx = q_idx[q_idx2]

    elif sampling_method == 'mc_dropout':
        # Assumes DNN 
        query_idx = mc_dropout(X_seed, y_seed, y_seed_true,
                            X_val, y_val, y_val_true, X_train_s, y_train, 
                            n_models = 20, 
                            variance=True, 
                            n_instances=num_instances,
                            dual=dual)
    elif sampling_method == 'ensemble':
        query_idx = ensemble(X_seed, y_seed, X_train, num_models = 5, n_instances=num_instances)
    else:
        print("This strategy is not implemented yet.")
        exit()
    
    del sub_model

    # Get the predictions from the Oracle (target model)
    X = X_train[query_idx]
    y = np.reshape(y_train[query_idx], (num_instances, ))
    if scaling_needed:
        X_s = X_train_s[query_idx] 
    
    if isAV:
        y_unlab = target_model(query_idx)
    else:
        y_unlab = target_model(X)
    y_unlab = np.reshape(y_unlab, (num_instances,))


    # Add the new data points and labels to the labeled pool
    if scaling_needed:
        X_seed = np.vstack((X_seed, X_s))
    else:    
        X_seed = np.vstack((X_seed, X))
    y_seed = np.concatenate((y_seed, y_unlab))
    y_seed_true = np.concatenate((y_seed_true, y))

    if model_type == 'DNN':
        sub_model = KerasAttacker(early_stopping=30, seed=seed)
        sub_model.train_model(X_seed, y_seed, X_val, y_val, num_epochs)
    elif model_type == 'dualDNN':
        sub_model = KerasDualAttacker(early_stopping=30, seed=seed)
        sub_model.train_model(X_seed, y_seed, y_seed_true, 
                        X_val, y_val, y_val_true, num_epochs)
    elif model_type == 'SVM':
        sub_model = SVMAttacker(seed=seed, max_iter=10000)
        sub_model.train_model(X_seed, y_seed)
    else:
        sub_model = LGBAttacker(seed=42)
        sub_model.train_model(X_seed, y_seed, X_val, y_val, boosting_rounds=1000, early_stopping=60)

    # Delete the used samples from the unlabeled pool
    X_train = np.delete(X_train, query_idx, axis=0)
    y_train = np.delete(y_train, query_idx, axis=0)
    if scaling_needed:
        X_train_s = np.delete(X_train_s, query_idx, axis=0)
    
    if dual:
        y_proba = sub_model(X_test, y_test)
    else:
        y_proba = sub_model(X_test) 
        
    num_samples = (idx+1)*num_instances+SEED_SIZE+VAL_SIZE
    scores["nums"].append(num_samples)
    scores = log_and_score(y_proba, y_pred_target, y_test, scores, fpr_target, logging=MLFLOW)


df = pd.DataFrame(scores)
if not os.path.exists(log_dir):
    os.makedirs(log_dir)

df.to_csv(os.path.join(log_dir, f"results_{target}_{seed}_{sampling_method}_{model_type}_{datetime.now().strftime('%H_%M_%d_%m_%Y')}.csv"), index=False)
print(scores)

sub_model.save_model(os.path.join(log_dir, f"model_{target}_{seed}_{sampling_method}_{model_type}_{datetime.now().strftime('%H_%M_%d_%m_%Y')}"))