BSidesLV -- your model isn't that special -- (3) MalwaResNet.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "env: CUDA_VISIBLE_DEVICES=0 # limit GPU usage, if any to this GPU\n"
     ]
    }
   ],
   "source": [
    "%env CUDA_VISIBLE_DEVICES=0 # limit GPU usage, if any to this GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from classifier import common\n",
    "import os\n",
    "labels = common.fetch_samples()\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "np.random.seed(123)\n",
    "y_train, y_test, sha256_train, sha256_test = train_test_split(\n",
    "    list(labels.values()), list(labels.keys()), test_size=1000)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <span style=\"color:#CC1111\">Malwa</span>ResNet\n",
    "\n",
    "Clearly, my model needs a heavy dose of special fanciness.  We'll turn to ResNet architecture!  This builds on the end-to-end model in the previous notebook, but stacks things higher and deeper, thanks to our hand-crafted residual cell.\n",
    "\n",
    "This model may require a hefty GPU.  I ran it successfully on a single TITAN X (Pascal).  If yours doesn't work, try enabling the `lite=True` options in create_model. \n",
    "\n",
    "You can find code that defines the MalwaResNet model architecture at [classifier/malwaresnet.py](classifier/malwaresnet.py)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# for this demo, will slurp in only the first 256K (2**18) bytes of the file\n",
    "max_file_length = int(2**18)\n",
    "file_chunks = 16  # break file into this many chunks\n",
    "file_chunk_size = max_file_length // file_chunks\n",
    "batch_size = 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "learning rate = 0.1\n",
      "Epoch 1/20\n",
      "24750/24750 [==============================] - 30939s - loss: 1.8194 - acc: 0.6434 - val_loss: 1.0777 - val_acc: 0.7020\n",
      "learning rate = 0.05\n",
      "Epoch 2/20\n",
      "24750/24750 [==============================] - 31095s - loss: 0.9919 - acc: 0.6874 - val_loss: 0.8848 - val_acc: 0.6950\n",
      "learning rate = 0.025\n",
      "Epoch 3/20\n",
      "24750/24750 [==============================] - 31509s - loss: 0.8041 - acc: 0.7172 - val_loss: 0.7376 - val_acc: 0.7740\n",
      "learning rate = 0.00625\n",
      "Epoch 5/20\n",
      "24750/24750 [==============================] - 31157s - loss: 0.7655 - acc: 0.7404 - val_loss: 0.6755 - val_acc: 0.8070\n",
      "learning rate = 0.003125\n",
      "Epoch 6/20\n",
      "24750/24750 [==============================] - 31124s - loss: 0.7669 - acc: 0.7300 - val_loss: 0.6939 - val_acc: 0.7580\n",
      "learning rate = 0.0015625\n",
      "Epoch 7/20\n",
      " 6591/24750 [======>.......................] - ETA: 22762s - loss: 0.7473 - acc: 0.7501"
     ]
    }
   ],
   "source": [
    "# That this is very long running cell, and \n",
    "# it may appear that the output is truncated before training completes\n",
    "\n",
    "# let's train this puppy\n",
    "from classifier import malwaresnet\n",
    "import math\n",
    "from keras.callbacks import LearningRateScheduler, EarlyStopping, ModelCheckpoint\n",
    "\n",
    "# create_model(input_shape, byte_embedding_size=2, lite=False)\n",
    "model_malwaresnet = malwaresnet.create_model(input_shape=(file_chunks, file_chunk_size), byte_embedding_size=2)\n",
    "train_generator = common.generator(list(zip(sha256_train, y_train)), batch_size, file_chunks, file_chunk_size)\n",
    "test_generator = common.generator(list(zip(sha256_test, y_test)), 1, file_chunks, file_chunk_size)\n",
    "training_history = model_malwaresnet.fit_generator(train_generator,\n",
    "                        steps_per_epoch=math.ceil(len(sha256_train) / batch_size),\n",
    "                        epochs=20,\n",
    "                        callbacks=[\n",
    "                            EarlyStopping( patience=1 ),\n",
    "                            ModelCheckpoint( 'malwaresnet.h5', save_best_only=True),\n",
    "                            LearningRateScheduler(lambda epoch: common.schedule(epoch, start=0.1, decay=0.5, every=1))],\n",
    "                        validation_data=test_generator,\n",
    "                        validation_steps=len(sha256_test))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Okay, no more snarkiness about impatient millenials from me.  **Each epoch** is taking 31000 s = 517 min = 8hrs 36 minutes.  But, we're going to remain optimistic that this is going to be awesome.  With a name like **<span style=\"color:#CC1111\">Malwa</span>ResNet**, how could it not be?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "** End-to-end convnet **\n",
      "ROC AUC = 0.8613780511243514\n",
      "threshold=0.8300181031227112: 0.08471074380165289 TP rate @ 0.009689922480620155 FP rate\n",
      "confusion matrix @ threshold:\n",
      "[[512   4]\n",
      " [443  41]]\n",
      "accuracy @ threshold = 0.553\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(0.8613780511243514,\n",
       " 0.8300181,\n",
       " 0.0096899224806201549,\n",
       " 0.084710743801652888,\n",
       " array([[512,   4],\n",
       "        [443,  41]]),\n",
       " 0.55300000000000005)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from keras.models import load_model\n",
    "# we'll load the \"best\" model (in this case, the penultimate model) that was saved \n",
    "# by our ModelCheckPoint callback\n",
    "model_malwaresnet = load_model('malwaresnet.h5')\n",
    "# we could load the \"best\" model, but in this case, the \"best\" model is the penultimate, and not much better\n",
    "# than the model we have in hand\n",
    "y_pred = []\n",
    "for sha256, lab in zip(sha256_test, y_test):\n",
    "    y_pred.append(\n",
    "        model_malwaresnet.predict_on_batch(\n",
    "            np.asarray([common.get_file_data(sha256, lab, max_file_length)]).reshape(\n",
    "                (-1, file_chunks, file_chunk_size))\n",
    "        )\n",
    "    )\n",
    "common.summarize_performance(np.asarray(y_pred).flatten(), y_test, \"End-to-end convnet\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# So much innovation!  So much GPU heat to the universe!  For this?!\n",
    "\n",
    "(┛◉Д◉)┛彡┻━┻\n",
    "\n",
    "For real, though.  Even if **<span style=\"color:#CC1111\">Malwa</span>ResNet** were a killer architecture for malware, it'd need to be trained on a lot more than a paltry 100K malicious/benign samples, and the optimization and architecture would need some revamping (for example, ResNet is trained over hundreds of epochs...on metric ton of relatively small images).  But as it is, this *special* model isn't really even close in performance to our *not special* multilayer perceptron.  (It's probable that if I let this train for another few weeks, we'd see some modest improvements...but not seeing it get near 0.99+ AUC.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Let's wrap this up\n",
    "\n",
    "Alright, the point of this exercise has been to demonstrate that\n",
    "1. You malware model isn't that special. \n",
    "2. Well, it might be special, but it probably isn't just your fancy architecture that makes it so. The first notebook in this series shows a *not special* model.\n",
    "3. Static malware detection isn't an \"easy win\" application for end-to-end deep learning models (e.g., ResNet) that were originally developed for image recognition.  Furthermore, our training set in these notebooks is much too small for the expressive end-to-end models with millions of parameters. \n",
    "4. End-to-end deep learning for static malware classification is a cool idea, but requires some work and perhaps more research to beat out *not special* models with great data, great features, and maticulously inspected/corrected labels.\n",
    "5. If you're looking for a quick win, go with the \"shallow\" model that *isn't that special*.  And do the work required to make its training data (and feature set) great.\n",
    "6. If you're a researcher looking to advance the state of the art in end-to-end deep leaning for static malware detection (no parsing), let's chat about creative ideas, layers and \"gadgets\" that might help get us there. \n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}