<gh_stars>0
"""
Copyright (c) 2021, FireEye, Inc.
Copyright (c) 2021 <NAME>

This script allows the user to train the models used in the paper.

To train the LightGBM and EmberNN models on EMBER:
`python train_model.py -m lightgbm -d ember`
`python train_model.py -m embernn -d ember`

To train the Random Forest model on Contagio PDFs:
`python train_model.py -m pdfrf -d ogcontagio`

To train the Linear SVM classifier on Drebin:
`python train_model.py -m linearsvm -d drebin`
"""

import random
import argparse
import os

import numpy as np
import tensorflow as tf

from mw_backdoor import constants
from mw_backdoor import data_utils
from mw_backdoor import model_utils


def train_eval(args):
    # Unpacking
    model_id = args['model']
    dataset = args['dataset']
    seed = args['seed']

    save_dir = args['save_dir']
    save_file = args['save_file']
    if not save_dir:
        save_dir = constants.SAVE_MODEL_DIR
    if not save_file:
        save_file = dataset + '_' + model_id
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)

    # Set random seeds
    random.seed(seed)
    np.random.seed(seed)
    tf.random.set_seed(seed)

    # Load data
    x_train, y_train, x_test, y_test = data_utils.load_dataset(dataset=dataset)
    if args["load_poison_dir"] is not None:
        x_train = np.load(os.path.join(args["load_poison_dir"], "watermarked_X.npy"))
        y_train = np.load(os.path.join(args["load_poison_dir"], "watermarked_y.npy"))
    if args["bagging"] is not None:
        idx = np.random.choice(np.arange(len(x_train)), args["bagging"], replace=True)
        x_train = x_train[idx]
        y_train = y_train[idx]
        # flip_mask = np.random.random(x_train.shape) < 0.3  # flip 30% of features
        # x_train = (1 - flip_mask) * x_train + flip_mask * (1 - x_train)
    elif args["partitioning"] is not None:
        assert args["partition_id"] is not None and 0 <= args["partition_id"] < args["partitioning"]
        idx = np.random.choice(np.arange(len(x_train)))
        num_per_partition = len(x_train) // args["partitioning"]
        x_train = x_train[idx][num_per_partition * args["partition_id"]:
                               num_per_partition * (args["partition_id"] + 1)]
        y_train = y_train[idx][num_per_partition * args["partition_id"]:
                               num_per_partition * (args["partition_id"] + 1)]

    print(
        'Dataset shapes:\n'
        '\tTrain x: {}\n'
        '\tTrain y: {}\n'
        '\tTest x: {}\n'
        '\tTest y: {}\n'.format(
            x_train.shape, y_train.shape, x_test.shape, y_test.shape
        )
    )

    if not args["eval"]:
        # Train model
        model = model_utils.train_model(
            model_id=model_id,
            x_train=x_train,
            y_train=y_train,
            args=args
        )

        # Save trained model
        model_utils.save_model(
            model_id=model_id,
            model=model,
            save_path=save_dir,
            file_name=save_file
        )
    else:
        assert os.path.exists(os.path.join(save_dir, "%s.h5" % save_file))
        model = model_utils.load_embernn(dataset, save_dir, save_file)

    # Evaluation
    raw_pred, pred = model_utils.evaluate_model(model, x_test, y_test)
    arguments["proba_delta"] = 0  # deprecated
    np.save(os.path.join(save_dir, save_file + "ori_test_pred"),
            np.array([pred == y_test, np.abs(raw_pred - 0.5) * 2 > arguments["proba_delta"]]))
    if args["load_poison_dir"] is not None:
        x_wm_test = np.load(os.path.join(args["load_poison_dir"], "watermarked_X_test.npy"))
        raw_pred_wm, pred_wm = model_utils.evaluate_model(model, x_wm_test, np.ones(len(x_wm_test)))
        np.save(os.path.join(save_dir, save_file + "wm_test_pred"),
                np.array([pred_wm == 1, np.abs(raw_pred_wm - 0.5) * 2 > arguments["proba_delta"]]))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "-m",
        "--model",
        default="lightgbm",
        choices=["lightgbm", "embernn", "pdfrf", "linearsvm"],
        help="model type"
    )
    parser.add_argument(
        "-d",
        "--dataset",
        default="ember",
        choices=["ember", "pdf", "ogcontagio", "drebin"],
        help="model type"
    )
    parser.add_argument(
        "--save_file",
        default='',
        type=str,
        help="file name of the saved model"
    )
    parser.add_argument(
        "--save_dir",
        default='',
        type=str,
        help="directory containing saved models"
    )
    parser.add_argument('--bagging', action='store', default=None, type=int,
                        help='Indicates whether and how many training data in the bagging')
    parser.add_argument('--partitioning', action='store', default=None, type=int,
                        help='Indicates whether and how many partitions we will create')
    parser.add_argument('--partition_id', action='store', default=None, type=int,
                        help='Indicates the partition id to train, partition_id is in [0, partitioning)')
    parser.add_argument('--early_stopping', action='store', default=None, type=int,
                        help='Indicates whether to use early stopping to stop the training')
    parser.add_argument('--discretize', action='store', default=None, type=int,
                        help='Number of bins for discretization')
    parser.add_argument(
        "--load_poison_dir",
        default=None,
        type=str,
        help="directory containing poisoned data"
    )
    parser.add_argument(
        "--eval",
        action='store_true',
        help="directly eval the model"
    )
    parser.add_argument(
        "--aug_noise",
        action='store_true',
        help="add sigma=0.1 Gaussian noise for data augmentation"
    )
    parser.add_argument(
        "--aug_flip",
        action='store_true',
        help="add flip noise (randomly choose 1% of the feature and set to arbitrary value) for data augmentation"
    )
    parser.add_argument("--epochs", type=int, default=10, help="training epochs")
    parser.add_argument("--batch_size", type=int, default=512, help="batch size")

    parser.add_argument("--seed", type=int, default=42, help="random seed")

    arguments = vars(parser.parse_args())
    train_eval(arguments)

"""
Copyright (c) 2021, FireEye, Inc.
Copyright (c) 2021 <NAME>

This script allows the user to train the models used in the paper.

To train the LightGBM and EmberNN models on EMBER:
`python train_model.py -m lightgbm -d ember`
`python train_model.py -m embernn -d ember`

To train the Random Forest model on Contagio PDFs:
`python train_model.py -m pdfrf -d ogcontagio`

To train the Linear SVM classifier on Drebin:
`python train_model.py -m linearsvm -d drebin`
"""

import random
import argparse
import os

import numpy as np
import tensorflow as tf

from mw_backdoor import constants
from mw_backdoor import data_utils
from mw_backdoor import model_utils


def train_eval(args):
    # Unpacking
    model_id = args['model']
    dataset = args['dataset']
    seed = args['seed']

    save_dir = args['save_dir']
    save_file = args['save_file']
    if not save_dir:
        save_dir = constants.SAVE_MODEL_DIR
    if not save_file:
        save_file = dataset + '_' + model_id
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)

    # Set random seeds
    random.seed(seed)
    np.random.seed(seed)
    tf.random.set_seed(seed)

    # Load data
    x_train, y_train, x_test, y_test = data_utils.load_dataset(dataset=dataset)
    if args["load_poison_dir"] is not None:
        x_train = np.load(os.path.join(args["load_poison_dir"], "watermarked_X.npy"))
        y_train = np.load(os.path.join(args["load_poison_dir"], "watermarked_y.npy"))
    if args["bagging"] is not None:
        idx = np.random.choice(np.arange(len(x_train)), args["bagging"], replace=True)
        x_train = x_train[idx]
        y_train = y_train[idx]
        # flip_mask = np.random.random(x_train.shape) < 0.3  # flip 30% of features
        # x_train = (1 - flip_mask) * x_train + flip_mask * (1 - x_train)
    elif args["partitioning"] is not None:
        assert args["partition_id"] is not None and 0 <= args["partition_id"] < args["partitioning"]
        idx = np.random.choice(np.arange(len(x_train)))
        num_per_partition = len(x_train) // args["partitioning"]
        x_train = x_train[idx][num_per_partition * args["partition_id"]:
                               num_per_partition * (args["partition_id"] + 1)]
        y_train = y_train[idx][num_per_partition * args["partition_id"]:
                               num_per_partition * (args["partition_id"] + 1)]

    print(
        'Dataset shapes:\n'
        '\tTrain x: {}\n'
        '\tTrain y: {}\n'
        '\tTest x: {}\n'
        '\tTest y: {}\n'.format(
            x_train.shape, y_train.shape, x_test.shape, y_test.shape
        )
    )

    if not args["eval"]:
        # Train model
        model = model_utils.train_model(
            model_id=model_id,
            x_train=x_train,
            y_train=y_train,
            args=args
        )

        # Save trained model
        model_utils.save_model(
            model_id=model_id,
            model=model,
            save_path=save_dir,
            file_name=save_file
        )
    else:
        assert os.path.exists(os.path.join(save_dir, "%s.h5" % save_file))
        model = model_utils.load_embernn(dataset, save_dir, save_file)

    # Evaluation
    raw_pred, pred = model_utils.evaluate_model(model, x_test, y_test)
    arguments["proba_delta"] = 0  # deprecated
    np.save(os.path.join(save_dir, save_file + "ori_test_pred"),
            np.array([pred == y_test, np.abs(raw_pred - 0.5) * 2 > arguments["proba_delta"]]))
    if args["load_poison_dir"] is not None:
        x_wm_test = np.load(os.path.join(args["load_poison_dir"], "watermarked_X_test.npy"))
        raw_pred_wm, pred_wm = model_utils.evaluate_model(model, x_wm_test, np.ones(len(x_wm_test)))
        np.save(os.path.join(save_dir, save_file + "wm_test_pred"),
                np.array([pred_wm == 1, np.abs(raw_pred_wm - 0.5) * 2 > arguments["proba_delta"]]))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "-m",
        "--model",
        default="lightgbm",
        choices=["lightgbm", "embernn", "pdfrf", "linearsvm"],
        help="model type"
    )
    parser.add_argument(
        "-d",
        "--dataset",
        default="ember",
        choices=["ember", "pdf", "ogcontagio", "drebin"],
        help="model type"
    )
    parser.add_argument(
        "--save_file",
        default='',
        type=str,
        help="file name of the saved model"
    )
    parser.add_argument(
        "--save_dir",
        default='',
        type=str,
        help="directory containing saved models"
    )
    parser.add_argument('--bagging', action='store', default=None, type=int,
                        help='Indicates whether and how many training data in the bagging')
    parser.add_argument('--partitioning', action='store', default=None, type=int,
                        help='Indicates whether and how many partitions we will create')
    parser.add_argument('--partition_id', action='store', default=None, type=int,
                        help='Indicates the partition id to train, partition_id is in [0, partitioning)')
    parser.add_argument('--early_stopping', action='store', default=None, type=int,
                        help='Indicates whether to use early stopping to stop the training')
    parser.add_argument('--discretize', action='store', default=None, type=int,
                        help='Number of bins for discretization')
    parser.add_argument(
        "--load_poison_dir",
        default=None,
        type=str,
        help="directory containing poisoned data"
    )
    parser.add_argument(
        "--eval",
        action='store_true',
        help="directly eval the model"
    )
    parser.add_argument(
        "--aug_noise",
        action='store_true',
        help="add sigma=0.1 Gaussian noise for data augmentation"
    )
    parser.add_argument(
        "--aug_flip",
        action='store_true',
        help="add flip noise (randomly choose 1% of the feature and set to arbitrary value) for data augmentation"
    )
    parser.add_argument("--epochs", type=int, default=10, help="training epochs")
    parser.add_argument("--batch_size", type=int, default=512, help="batch size")

    parser.add_argument("--seed", type=int, default=42, help="random seed")

    arguments = vars(parser.parse_args())
    train_eval(arguments)
Copyright (c) 2021, FireEye, Inc. Copyright (c) 2021 <NAME>  This script allows the user to train the models used in the paper.  To train the LightGBM and EmberNN models on EMBER: `python train_model.py -m lightgbm -d ember` `python train_model.py -m embernn -d ember`  To train the Random Forest model on Contagio PDFs: `python train_model.py -m pdfrf -d ogcontagio`  To train the Linear SVM classifier on Drebin: `python train_model.py -m linearsvm -d drebin` # Unpacking # Set random seeds # Load data # flip_mask = np.random.random(x_train.shape) < 0.3  # flip 30% of features # x_train = (1 - flip_mask) * x_train + flip_mask * (1 - x_train) # Train model # Save trained model # Evaluation # deprecated