models.py

# coding: utf-8

# In[1]:


import tensorflow as tf
import numpy as np
from sklearn.metrics import confusion_matrix
from sklearn.model_selection import train_test_split
from sklearn.utils import shuffle
import time
from datetime import timedelta
import math
import _pickle as cPickle
import os


# In[2]:


# # Embedding Layer
# num_rows = 218
# num_cols = 8

# Convolutional Layer
filter_size = 8 
num_filters = 64

# Max Pooling Layer


# Fully-connected Layer
fc_size = 16


# In[3]:


def load_data(folder):
    x = []
    y = []
    for path in os.listdir(folder):
        input = cPickle.load(open(os.path.join(folder, path), 'rb'))
        x.append(input['x'])
        y.append(input['y'])
    return x, y


# In[4]:


x, y = load_data('./Pickles')
# x2, y2 = load_data('../opseq_benign/')

# x = x1 + x2
# y = y1 + y2


# In[5]:


x, y = shuffle(x, y, random_state=137)
x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=137)


# In[6]:


# Each opcode sequence in one-hot form is of dimensions N x 256
op_rows = 100000
op_cols = 1

# Opcode sequences are stored in one-dimensional arrays of this length
op_size_flat = op_rows * op_cols

# Tuple with height and width of opcode sequences used to reshape arrays
op_shape = (op_rows, op_cols)

# Number of channels for the input: 1
num_channels = 1

# Number of classes
num_classes = 2


# In[7]:


def new_weights(shape):
	return tf.Variable(tf.truncated_normal(shape, stddev=0.05))

def new_biases(length):
	return tf.Variable(tf.constant(0.05, shape=[length]))


# In[8]:


def new_conv_layer(input,              	# The input layer
                   num_input_channels, 	# Number of channels in input layer
                   filter_size,        	# Width and height of each filter
                   num_filters,        	# Number of filters
                   use_pooling=False, name="conv"):  # Max Pooling not done
    
    with tf.name_scope(name):
        # Shape of the filter-weights for the convolution.
        shape = [filter_size, filter_size, num_input_channels, num_filters]

        # Create new weights aka. filters with the given shape.
        weights = new_weights(shape=shape)

        # Create new biases, one for each filter.
        biases = new_biases(length=num_filters)

        # Create the TensorFlow operation for convolution.
        layer = tf.nn.conv2d(input=input,
                             filter=weights,
                             strides=[1, 1, 1, 1],
                             padding='VALID')

        # Add the biases to the results of the convolution.
        layer += biases

        # Use pooling to down-sample the image resolution?
        if use_pooling:
            # This is 2x2 max-pooling, which means that we
            # consider 2x2 windows and select the largest value
            # in each window. Then we move 2 pixels to the next window.
            layer = tf.nn.max_pool(value=layer,
                                   ksize=[1, 2, 2, 1],
                                   strides=[1, 2, 2, 1],
                                   padding='VALID')

        # Rectified Linear Unit (ReLU).
        layer = tf.nn.relu(layer)
        tf.summary.histogram(name + "/weights", weights)
        tf.summary.histogram(name + "/biases", biases)
        tf.summary.histogram(name+ "/activations", layer)
        return layer, weights


# In[9]:


def new_fc_layer(input,          # The previous layer
                 num_inputs,     # Number of inputs from previous layer
                 num_outputs,    # Number of outputs
                 use_relu=True, name="fc"): # Use Rectified Linear Unit (ReLU)?

    with tf.name_scope(name):
        # Create new weights and biases.
        weights = new_weights(shape=[num_inputs, num_outputs])
    #     weights = tf.expand_dims(weights,0)
        biases = new_biases(length=num_outputs)

        # Calculate the layer as the matrix multiplication of
        # the input and weights, and then add the bias-values.
        layer = tf.matmul(input, weights) + biases

        # Use ReLU?
        if use_relu:
            layer = tf.nn.relu(layer)
        tf.summary.histogram(name + "/weights", weights)
        tf.summary.histogram(name + "/biases", biases)
        tf.summary.histogram(name+ "/activations", layer)
        return layer


# In[10]:


#Placeholder for 1d input in ints
x = tf.placeholder(tf.int32, shape=[None, op_rows])
#Placeholder for one hot
# x_one_hot = tf.placeholder(tf.float32, shape=[None, op_rows, op_cols])
# Placeholder for input
x_ints = tf.placeholder(tf.int32, shape=[None, op_size_flat], name='x')
# Converting into a four dimensional vector
x_inp = tf.reshape(x, [-1, op_rows, op_cols, num_channels])
# Placeholder for true labels
y_true = tf.placeholder(tf.float32, shape=[None, num_classes], name='y_true')
# Placeholder variable for class number
y_true_cls = tf.argmax(y_true, axis=1)


# In[11]:


embeddings = tf.Variable(
    tf.random_uniform(
    [218,8],
    minval=-1,
    maxval=1,
    dtype=tf.float32,
    seed=None,
    name=None
))
embed = tf.nn.embedding_lookup(embeddings, x_ints)
embed = tf.reshape(embed, shape=[-1, 100000, 8, 1])
print(embed)


# In[12]:


layer_conv, weights_conv =     new_conv_layer(input=embed,
                   num_input_channels=num_channels,
                   filter_size=filter_size,
                   num_filters=num_filters,
                   use_pooling=False, name="conv1")


# In[13]:


print(layer_conv)


# In[14]:


max_layer = tf.reduce_max(
    layer_conv,
    axis=1,
#     keepdims=False,
    name=None,
    reduction_indices=None,
    keep_dims=False
)
max_layer = tf.reshape(max_layer, shape=[-1,64])
# print max_layer


# In[15]:


layer_fc = new_fc_layer(input=max_layer,
                         num_inputs=64,
                         num_outputs=2,
                         use_relu=False, name="fc1")
# print layer_fc


# In[16]:


y_pred = tf.nn.softmax(layer_fc)
# print y_pred


# In[17]:


y_pred_cls = tf.argmax(y_pred, axis=1)


# In[18]:


with tf.name_scope("xent"):
    cross_entropy = tf.nn.softmax_cross_entropy_with_logits(logits=layer_fc, labels=y_true)
    cost = tf.reduce_mean(cross_entropy)
    tf.summary.scalar("xent", cost)


# In[19]:


optimizer = tf.train.AdamOptimizer(learning_rate=1e-4).minimize(cost)


# In[20]:


with tf.name_scope("accuracy"):
    correct_prediction = tf.equal(y_pred_cls, y_true_cls)
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
    tf.summary.scalar("accuracy", accuracy)


# In[21]:


writer = tf.summary.FileWriter("temp/1", graph=tf.get_default_graph())
merged_summary = tf.summary.merge_all()


# In[22]:


session = tf.Session()


# In[23]:


session.run(tf.global_variables_initializer())


# In[24]:


train_batch_size = 1


# In[25]:


def optimize(num_epochs, batch_size):
    counter=0
    for epoch in range(num_epochs):
        for step in range(int(len(x_train)/batch_size)):
            batch_x = x_train[step*batch_size:(step+1)*batch_size]
            batch_y = y_train[step*batch_size:(step+1)*batch_size]
            b = np.zeros((batch_size,2))
            b[np.arange(batch_size),np.asarray(batch_y)] = 1
            batch_y = np.float32(b)
            batch_x = np.float32(batch_x)
            assert(batch_x.shape[0]==batch_y.shape[0])
            feed_dict_train = {x_ints: batch_x,
                               y_true: batch_y}
            session.run(optimizer, feed_dict=feed_dict_train)
            
            acc = session.run(accuracy, feed_dict=feed_dict_train)
            summary = session.run(merged_summary, feed_dict=feed_dict_train)
            writer.add_summary(summary, counter)
            
            msg = "Optimization Iteration: {0:>6}, Training Accuracy: {1:>6.1%}"

            # Print it.
            print(msg.format(counter + 1, acc))
            counter = counter + 1


# In[26]:


optimize(num_epochs=3, batch_size=train_batch_size)


# In[27]:


# saver = tf.train.Saver()
# saver.restore(sess, 'tmp/1/model.cpkt')


# In[29]:


saver = tf.train.Saver()
saver.save(session, 'Desktop/model.cpkt')


# In[30]:


b = np.zeros((584,2))
b[np.arange(584),np.asarray(y_test)] = 1
y_test = np.float32(b)
feed_dict_test = {x_ints: x_test,
                y_true: y_test}


# In[37]:


#accuracy test data

batch_size = 50
avg_accuracy = []
for i in range(int(584 / batch_size)):
    batch_x = x_test[i * batch_size:(i + 1) * batch_size]
    batch_y = y_test[i * batch_size:(i + 1) * batch_size]
    feed_dict_test = {x_ints: batch_x, y_true: batch_y}
    acc = session.run(accuracy, feed_dict=feed_dict_test)
    avg_accuracy.append(acc)
    print("Batch: %d -- Accuracy: %g" %(i, acc))


# In[ ]:


print("Average Accuracy: %g", np.mean(avg_accuracy))


# In[ ]:


# session.close()