dl-malware-detection/mainRF.py at main · willmacd/dl-malware-detection · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
import numpy as np

from DataProcessing.LoadData import dat_to_train_test, to_tf_dataset, \
    rm_unlabelled_samples, to_batch_dataset, split_data_labels
from Model.RFMalwareDetection import MalwareDetectionRF

FILTERED_DATASET_SIZES = 600000

BATCH_SIZE = 1000
EPOCHS = 125

# CHECK: Ensure that this is the correct path to the dataset
DATA_DIR = './Data/dat/'

if __name__ == '__main__':

    x_train, y_train, x_test, y_test = dat_to_train_test(DATA_DIR)

    '''# Normalize the data using robust scaler
    print("Normalizing data...")
    x_train_scaled = normalize_data(x_train)
    x_test_scaled = normalize_data(x_test)
    print("Data normalization complete...")

    # Apply PCA dimensionality reduction
    print("Computing PCA for dimensionality reduction...")
    x_train_pca, x_test_pca = dataset_pca_reduction(x_train_scaled, x_test_scaled)
    print("PCA dimensionality reduction complete...")

    # Number of components kept by PCA
    num_components_pca = len(x_train_pca[1])
    print(num_components_pca)'''

    unfiltered_train_ds = to_tf_dataset(x_train, y_train)   # _pca
    unfiltered_test_ds = to_tf_dataset(x_test, y_test)  # _pca

    # Filter out the data with label '-1' (unlabeled)
    filtered_train_ds = rm_unlabelled_samples(unfiltered_train_ds)
    filtered_test_ds = rm_unlabelled_samples(unfiltered_test_ds)

    train_ds = filtered_train_ds.take(int(0.85 * FILTERED_DATASET_SIZES))
    val_ds = filtered_train_ds.skip(int(0.85 * FILTERED_DATASET_SIZES))

    train_ds = to_batch_dataset(filtered_train_ds, BATCH_SIZE)
    val_ds = to_batch_dataset(val_ds, BATCH_SIZE)
    test_ds = to_batch_dataset(filtered_test_ds, BATCH_SIZE)

    train_data, train_labels = split_data_labels(train_ds, 10000)
    test_data, test_labels = split_data_labels(test_ds, 1000)

    RandomForest = MalwareDetectionRF(num_trees=1000, verbose=1)
    print("Training Random Forest...")
    RandomForest.train(train_data, train_labels)
    print("Random Forest successfully trained...")
    print(RandomForest.evaluate(train_data, train_labels))

    print(RandomForest.evaluate(test_data, test_labels))