wip.py

"""main executing script
currently this script will train the model based on the configurations of the variables below
default 50 epochs
and analyse the accuracy and performance in a metric that i don't yet understand
"""

import os
import cv2
import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import matplotlib.pyplot as plt
from sklearn.metrics import classification_report, confusion_matrix
from keras.callbacks import ModelCheckpoint
from sklearn.model_selection import train_test_split

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
# setting the path for the dataset and model
path = "./exeimg_dataset"
model_path = "./model"

# Defining hyperparameters
num_classes = 2
input_shape = (299, 299, 3)
learning_rate = 0.001
weight_decay = 0.0001
num_epochs = 5
batch_size = 32
batch_testsize = 28
image_size = 128  # We will resize the input images to this size
patch_size = 8  # Size of the patches to be extracted from the input images
num_patches = (image_size // patch_size) ** 2
projection_dim = 64
num_heads = 2  # number of attention heads
transformer_units = [
    projection_dim * 2,
    projection_dim,
]  # Size of the transformer layers
transformer_layers = 1  # number of transformer layers
mlp_head_units = [2048, 1024]  # Size of the dense layers of the final classifier
key_dim = projection_dim // num_heads

# Creating the datasets
print("creating datasets")
X = []
y = []
files = ["benignware", "malware"]
label_val = 0

for f in files:
    cpath = os.path.join(path, f)
    print(f"cpath: {cpath}")
    for img in os.listdir(cpath):
        image_array = cv2.imread(os.path.join(cpath, img), cv2.IMREAD_COLOR)
        X.append(image_array)
        y.append(label_val)
    label_val = 1

X = np.asarray(X)
y = np.asarray(y)

# Set aside 20% for test and the remaining for train and validation sets
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, shuffle=True, random_state=8
)

# Set aside 10% for validation set and remaining for training set
X_train, X_val, y_train, y_val = train_test_split(
    X_train, y_train, test_size=0.1, random_state=8
)

# Checking the shape of the datasets
print(f"X_train: {X_train.shape} - y_train: {y_train.shape}")
print(f"X_val: {X_val.shape} - y_val: {y_val.shape}")
print(f"X_test: {X_test.shape} - y_test: {y_test.shape}")

# Calculating the steps per epoch
STEP_SIZE_TRAIN = int(np.ceil(X_train.shape[0] / batch_size))
STEP_SIZE_VALID = int(np.ceil(X_val.shape[0] / batch_size))
STEP_SIZE_TEST = int(np.ceil(X_test.shape[0] / batch_testsize))

print(f"train: {STEP_SIZE_TRAIN}, valid: {STEP_SIZE_VALID}, test: {STEP_SIZE_TEST}")

# Performing Data augmentation
data_augmentation = keras.Sequential(
    [
        layers.Normalization(),
        layers.Resizing(image_size, image_size),
    ],
    name="data_augmentation",
)
print("data augmentation done")

# Computing the mean and variance of the training data for normalization.
data_augmentation.layers[0].adapt(X_train)
print("adapting data augmentation")


# Building MLP network
def mlp(x, hidden_units, dropout_rate):
    print("running mlp function")
    for units in hidden_units:
        x = layers.Dense(units, activation=tf.nn.gelu)(x)
        x = layers.Dropout(dropout_rate)(x)
    return x


# Creating patches using tensorflow's extract_patches module
class Patches(layers.Layer):
    def __init__(self, patch_size):
        super(Patches, self).__init__()
        self.patch_size = patch_size

    def get_config(self):
        config = {
            "patch_size": self.patch_size,
        }
        base_config = super().get_config()
        return dict(list(base_config.items()) + list(config.items()))

    def call(self, images):
        batch_size = tf.shape(images)[0]
        patches = tf.image.extract_patches(
            images=images,
            sizes=[1, self.patch_size, self.patch_size, 1],
            strides=[1, self.patch_size, self.patch_size, 1],
            rates=[1, 1, 1, 1],
            padding="VALID",
        )
        patch_dims = patches.shape[-1]
        patches = tf.reshape(patches, [batch_size, -1, patch_dims])
        return patches


# Encoding the patches and adding postion embeddings
class PatchEncoder(layers.Layer):
    def __init__(self, num_patches, projection_dim):
        super(PatchEncoder, self).__init__()
        self.num_patches = num_patches
        self.projection = layers.Dense(units=projection_dim)
        self.position_embedding = layers.Embedding(
            input_dim=num_patches, output_dim=projection_dim
        )

    def get_config(self):
        config = {
            "num_patches": self.num_patches,
        }
        base_config = super().get_config()
        return dict(list(base_config.items()) + list(config.items()))

    def call(self, patch):
        positions = tf.range(start=0, limit=self.num_patches, delta=1)
        encoded = self.projection(patch) + self.position_embedding(positions)
        return encoded


# Building Vision Transformer
def create_vit_classifier():
    inputs = layers.Input(shape=input_shape)
    print(f"shape of inputs: {inputs.shape}")

    # Augment data
    augmented = data_augmentation(inputs)
    print(f"shape of augmented: {augmented.shape}")

    # Create patches
    patches = Patches(patch_size)(augmented)
    print(f"shape of patches: {patches.shape}")

    # Encode patches
    encoded_patches = PatchEncoder(num_patches, projection_dim)(patches)
    print(f"shape of encoded_patches: {encoded_patches.shape}")

    # Create multiple layers of the Transformer block
    for _ in range(transformer_layers):
        # Layer normalization 1
        x1 = layers.LayerNormalization(epsilon=1e-6)(encoded_patches)
        # Create a multi-head attention layer
        attention_output = layers.MultiHeadAttention(
            num_heads=num_heads,
            key_dim=key_dim,
            value_dim=key_dim,
            use_bias=False,
            dropout=0.1,
        )(x1, x1)
        print(f"shape of attention_output: {attention_output.shape}")
        # Skip connection 1
        x2 = layers.Add()([attention_output, encoded_patches])
        # Layer normalization 2
        x3 = layers.LayerNormalization(epsilon=1e-6)(x2)
        # MLP
        x3 = mlp(x3, hidden_units=transformer_units, dropout_rate=0.1)
        # Skip connection 2
        encoded_patches = layers.Add()([x3, x2])

    # Create a [batch_size, projection_dim] tensor
    representation = layers.LayerNormalization(epsilon=1e-6)(encoded_patches)
    representation = layers.Flatten()(representation)
    representation = layers.Dropout(0.5)(representation)
    # Add MLP
    features = mlp(representation, hidden_units=mlp_head_units, dropout_rate=0.5)
    print(f"shape of mlp output: {features.shape}")
    # Classify outputs
    logits = layers.Dense(num_classes)(features)
    print(f"shape of logits: {logits.shape}")
    # Create the Keras model
    model = keras.Model(inputs=inputs, outputs=logits)

    return model


def run_experiment(model):
    # Compiling the model
    print("optimizer")
    optimizer = tf.keras.optimizers.AdamW(
        learning_rate=learning_rate, weight_decay=weight_decay
    )

    model.compile(
        optimizer=optimizer,
        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=[
            keras.metrics.SparseCategoricalAccuracy(name="accuracy"),
            keras.metrics.SparseTopKCategoricalAccuracy(5, name="top-5-accuracy"),
        ],
    )

    # Save the Keras model weights at some frequency (here, when the best validation accuracy is achieved)
    checkpoint_filepath = "./tmp/checkpoint"
    checkpoint_callback = ModelCheckpoint(
        checkpoint_filepath,
        monitor="val_accuracy",
        mode="max",
        save_best_only=True,
        save_weights_only=True,
    )

    # Training the model
    history = model.fit(
        x=X_train,
        y=y_train,
        steps_per_epoch=STEP_SIZE_TRAIN,
        batch_size=batch_size,
        epochs=num_epochs,
        validation_data=(X_val, y_val),
        validation_steps=STEP_SIZE_VALID,
        callbacks=[checkpoint_callback],
        shuffle=True,
    )

    # Save the model
    model.save(os.path.join(model_path, "vit-malware-det-model.h5"))

    # Checking the accuracy
    model.load_weights(checkpoint_filepath)
    _, accuracy, top_5_accuracy = model.evaluate(
        X_test, y_test, batch_size=batch_testsize
    )
    print(f"test accuracy: {round(accuracy * 100, 2)}%")
    print(f"test top 5 accuracy: {round(top_5_accuracy * 100, 2)}%")

    # Plotting the accuracies
    plt.plot(history.history["accuracy"], label="train acc")
    plt.plot(history.history["val_accuracy"], label="val acc")
    plt.title("model accuracy")
    plt.ylabel("accuracy")
    plt.xlabel("epoch")
    plt.legend()
    plt.savefig(os.path.join(model_path, "accuracy.png"))
    plt.clf()

    # Plotting the loss
    plt.plot(history.history["loss"], label="train loss")
    plt.plot(history.history["val_loss"], label="val loss")
    plt.title("model loss")
    plt.ylabel("loss")
    plt.xlabel("epoch")
    plt.legend()
    plt.savefig(os.path.join(model_path, "loss.png"))

    # Making predictions
    y_pred = model.predict(X_test, steps=STEP_SIZE_TEST)
    y_pred = np.argmax(y_pred, axis=1)

    labels = {"benignware": 0, "malware": 1}
    print(f"labels and their corresponding encodings: {labels}")
    labels = dict((v, k) for k, v in labels.items())
    predictions = [labels[k] for k in y_pred]
    print(f"predictions: {y_pred}")
    print(f"predictions: {predictions}")

    # Get the classification report
    print(classification_report(y_pred, y_test))

    # Get the confusion matrix
    print(confusion_matrix(y_pred, y_test))

    return history


vit_classifier = create_vit_classifier()
vit_classifier.summary()
history = run_experiment(vit_classifier)