NuralNetwork.py

#!/usr/bin/env python
# coding: utf-8

# In[7]:


import tensorflow as tf
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
import os 
from tqdm import tqdm


# In[8]:


os.environ['CUDA_VISIBLE_DEVICES'] = ''


# In[9]:


Agreement_csv = ['bittorrent.csv',
                 'dns.csv',
                 'ftp.csv',
                 'httphttps.csv',
                 'pop3.csv',
                 'smtp.csv',
                 'ssh.csv',
                 'telnet.csv']


# In[10]:


def strip(data):
    columns = data.columns
    new_columns = []
    for i in range(len(columns)):
        new_columns.append(columns[i].strip())
    return new_columns


# In[11]:


path = '../Data'
for root, dirs, files in os.walk(path):  
    if files==[]:
        break
    csv = files


# In[12]:


csv = [i for i in csv if i not in Agreement_csv]


# In[13]:


csv


# In[14]:


data = pd.DataFrame()
for single_csv in csv:
    if 'csv' in single_csv:
        temp = pd.read_csv(os.path.join(path,single_csv))
    else:
        temp = pd.read_excel(os.path.join(path,single_csv))
    data = pd.concat([temp,data],axis=0)
data.columns = strip(data)


# In[15]:


Label_encoder = {data['Label'].unique()[i]:i for i in range(len(data['Label'].unique()))}


# In[16]:


data['Label']=data['Label'].apply(lambda x:Label_encoder[x])


# In[17]:


positive_len = len(data[data.Label!=1])
negative = data[data['Label'] == 1].sample(n = positive_len)
data_balance = pd.concat([negative,data[data.Label!=1]]).reset_index()
input_data = data_balance[data_balance.columns[:-1]]
input_labels = data_balance['Label']


# In[18]:


def delete_inf(data,label):
    input_data_ = np.array(data)
    input_labels_ = np.array(label)
    inf = np.argwhere(input_data_==np.inf)[:,0]
    inf = list(set(inf))
    input_data_ = np.delete(input_data_,inf,axis=0)
    input_labels_ = np.delete(input_labels_,inf)
    return input_data_,input_labels_


# In[19]:


input_data, input_labels = delete_inf(input_data, input_labels)


# In[20]:


input_data.shape


# In[21]:


train_data,test_data,train_labels,test_labels = train_test_split(input_data,input_labels,train_size=0.8, test_size=0.2, random_state=41)


# In[22]:


class NerualNetwork(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.layer1 = tf.keras.layers.Dense(units=128, activation=tf.nn.sigmoid)
        self.layer2 = tf.keras.layers.Dense(units=64, activation=tf.nn.relu)
        self.layer3 = tf.keras.layers.Dense(units=32, activation=tf.nn.relu)
        self.layer4 = tf.keras.layers.Dense(units=6, activation=tf.nn.relu)
        
    @tf.function
    def call(self,inputs):
        x = self.layer1(inputs)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        output = tf.nn.softmax(x)
        return output


# In[23]:


network = NerualNetwork()
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-5)


# In[24]:


train_data_keras = np.array(train_data)
test_data_keras = np.array(test_data)
y_train_keras = np.array(train_labels)
y_test_keras = np.array(test_labels)
batch_size=512


# In[25]:


def onehot_embedding(label):
    onehot_dict = {0:[1,0,0,0,0,0],1:[0,1,0,0,0,0],2:[0,0,1,0,0,0],3:[0,0,0,1,0,0],4:[0,0,0,0,1,0],5:[0,0,0,0,0,1]}
    one_hot = []
    for i in label:
        one_hot.append(onehot_dict[i])
    return np.array(one_hot)


# In[26]:


valid_data = (test_data_keras, onehot_embedding(y_test_keras))


# In[18]:


network.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
    loss=tf.keras.losses.CategoricalCrossentropy(),
    metrics=[tf.keras.metrics.CategoricalAccuracy()]
)
network.fit(train_data_keras, onehot_embedding(y_train_keras), epochs=2000, batch_size=batch_size,          verbose=1,validation_data=valid_data,callbacks=None)


# In[ ]:


tf.saved_model.save(network, "Model/nuralnetwork")


# In[28]:


model = tf.saved_model.load("../Model/nuralnetwork")


# In[29]:


def TopkAccuracy(label,input_data,input_label,k):
    pred_logits = model.call(tf.convert_to_tensor(test_data_keras,dtype=tf.float32))
    topk = np.argsort(np.array(pred_logits),axis=1)[:,-k:]
    nums=0
    cnt=0
    for index,values in enumerate(input_label):
        if values==label:
            nums+=1
            if values in topk[index]:
                cnt+=1
    return cnt/nums


# In[30]:


def CategoricalAccuracy(label,intput_data,input_label):
    acc = 0
    num = 0
    pred_logits = model.call(tf.convert_to_tensor(test_data_keras,dtype=tf.float32))
    predict = tf.argmax(pred_logits,axis=1)
    for index,values in enumerate(input_label):
        if values==label:
            num+=1
            if predict[index]==values:
                acc+=1
    return acc/num


# In[31]:


import matplotlib.pyplot as plt


# In[32]:


acc = []
for label in range(6):
    acc.append(TopkAccuracy(label,test_data_keras,y_test_keras,1))
acc 


# In[33]:


label = ['worm',
 'BENIGN',
 'XSS',
 'Sql Injection',
 'Trojan',
 'Overflow']


# In[34]:


fig,ax = plt.subplots(figsize=(10,5))
b1 = ax.bar(label[0],acc[0])
b2 = ax.bar(label[1],acc[1])
b3 = ax.bar(label[2],acc[2])
b4 = ax.bar(label[3],acc[3])
b5 = ax.bar(label[4],acc[4])
b6 = ax.bar(label[5],acc[5])
for b in b1+b2+b3+b4+b5+b6:
    h = b.get_height()
    ax.text(b.get_x()+b.get_width()/2,h,'%f'%h,ha='center',va='bottom')
plt.xlabel('Label')
plt.ylabel('accuracy')


# In[ ]: