Model-2_Balanced_Dataset.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Deep Graph Convolutional Neural Network (DGCNN)\n",
    "\n",
    "This code is part our research on malware detection and classification using Deep Learning and Deep Graph Convolutional Neural Networks.\n",
    "\n",
    "For more information or citation, please refer to our research paper:\n",
    "\n",
    "\"Oliveira, Angelo; Sassi, Renato José (2019): Behavioral Malware Detection Using Deep Graph Convolutional Neural Networks. TechRxiv. Preprint.\" at https://doi.org/10.36227/techrxiv.10043099.v1\n",
    "\n",
    "For the dataset, please refer to our repository:\n",
    "\n",
    "https://ieee-dataport.org/open-access/malware-analysis-datasets-api-call-sequences\n",
    "\n",
    "#### Model-2, Balanced Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "SEED = 137\n",
    "np.random.seed(SEED)\n",
    "\n",
    "import time\n",
    "from collections import Counter\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "from sklearn.model_selection import train_test_split, GridSearchCV, StratifiedKFold\n",
    "from sklearn.metrics import accuracy_score, recall_score, confusion_matrix\n",
    "from sklearn.metrics import precision_score, f1_score, roc_auc_score\n",
    "\n",
    "from imblearn.under_sampling import RandomUnderSampler\n",
    "from imblearn.pipeline import Pipeline\n",
    "\n",
    "import torch\n",
    "torch.manual_seed(SEED)\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from skorch.classifier import NeuralNetBinaryClassifier"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data Preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>hash</th>\n",
       "      <th>t_0</th>\n",
       "      <th>t_1</th>\n",
       "      <th>t_2</th>\n",
       "      <th>t_3</th>\n",
       "      <th>t_4</th>\n",
       "      <th>t_5</th>\n",
       "      <th>t_6</th>\n",
       "      <th>t_7</th>\n",
       "      <th>t_8</th>\n",
       "      <th>...</th>\n",
       "      <th>t_91</th>\n",
       "      <th>t_92</th>\n",
       "      <th>t_93</th>\n",
       "      <th>t_94</th>\n",
       "      <th>t_95</th>\n",
       "      <th>t_96</th>\n",
       "      <th>t_97</th>\n",
       "      <th>t_98</th>\n",
       "      <th>t_99</th>\n",
       "      <th>malware</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>071e8c3f8922e186e57548cd4c703a5d</td>\n",
       "      <td>112</td>\n",
       "      <td>274</td>\n",
       "      <td>158</td>\n",
       "      <td>215</td>\n",
       "      <td>274</td>\n",
       "      <td>158</td>\n",
       "      <td>215</td>\n",
       "      <td>298</td>\n",
       "      <td>76</td>\n",
       "      <td>...</td>\n",
       "      <td>71</td>\n",
       "      <td>297</td>\n",
       "      <td>135</td>\n",
       "      <td>171</td>\n",
       "      <td>215</td>\n",
       "      <td>35</td>\n",
       "      <td>208</td>\n",
       "      <td>56</td>\n",
       "      <td>71</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>33f8e6d08a6aae939f25a8e0d63dd523</td>\n",
       "      <td>82</td>\n",
       "      <td>208</td>\n",
       "      <td>187</td>\n",
       "      <td>208</td>\n",
       "      <td>172</td>\n",
       "      <td>117</td>\n",
       "      <td>172</td>\n",
       "      <td>117</td>\n",
       "      <td>172</td>\n",
       "      <td>...</td>\n",
       "      <td>81</td>\n",
       "      <td>240</td>\n",
       "      <td>117</td>\n",
       "      <td>71</td>\n",
       "      <td>297</td>\n",
       "      <td>135</td>\n",
       "      <td>171</td>\n",
       "      <td>215</td>\n",
       "      <td>35</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>b68abd064e975e1c6d5f25e748663076</td>\n",
       "      <td>16</td>\n",
       "      <td>110</td>\n",
       "      <td>240</td>\n",
       "      <td>117</td>\n",
       "      <td>240</td>\n",
       "      <td>117</td>\n",
       "      <td>240</td>\n",
       "      <td>117</td>\n",
       "      <td>240</td>\n",
       "      <td>...</td>\n",
       "      <td>65</td>\n",
       "      <td>112</td>\n",
       "      <td>123</td>\n",
       "      <td>65</td>\n",
       "      <td>112</td>\n",
       "      <td>123</td>\n",
       "      <td>65</td>\n",
       "      <td>113</td>\n",
       "      <td>112</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>72049be7bd30ea61297ea624ae198067</td>\n",
       "      <td>82</td>\n",
       "      <td>208</td>\n",
       "      <td>187</td>\n",
       "      <td>208</td>\n",
       "      <td>172</td>\n",
       "      <td>117</td>\n",
       "      <td>172</td>\n",
       "      <td>117</td>\n",
       "      <td>172</td>\n",
       "      <td>...</td>\n",
       "      <td>208</td>\n",
       "      <td>302</td>\n",
       "      <td>208</td>\n",
       "      <td>302</td>\n",
       "      <td>187</td>\n",
       "      <td>208</td>\n",
       "      <td>302</td>\n",
       "      <td>228</td>\n",
       "      <td>302</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>c9b3700a77facf29172f32df6bc77f48</td>\n",
       "      <td>82</td>\n",
       "      <td>240</td>\n",
       "      <td>117</td>\n",
       "      <td>240</td>\n",
       "      <td>117</td>\n",
       "      <td>240</td>\n",
       "      <td>117</td>\n",
       "      <td>240</td>\n",
       "      <td>117</td>\n",
       "      <td>...</td>\n",
       "      <td>209</td>\n",
       "      <td>260</td>\n",
       "      <td>40</td>\n",
       "      <td>209</td>\n",
       "      <td>260</td>\n",
       "      <td>141</td>\n",
       "      <td>260</td>\n",
       "      <td>141</td>\n",
       "      <td>260</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 102 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                               hash  t_0  t_1  t_2  t_3  t_4  t_5  t_6  t_7  \\\n",
       "0  071e8c3f8922e186e57548cd4c703a5d  112  274  158  215  274  158  215  298   \n",
       "1  33f8e6d08a6aae939f25a8e0d63dd523   82  208  187  208  172  117  172  117   \n",
       "2  b68abd064e975e1c6d5f25e748663076   16  110  240  117  240  117  240  117   \n",
       "3  72049be7bd30ea61297ea624ae198067   82  208  187  208  172  117  172  117   \n",
       "4  c9b3700a77facf29172f32df6bc77f48   82  240  117  240  117  240  117  240   \n",
       "\n",
       "   t_8  ...  t_91  t_92  t_93  t_94  t_95  t_96  t_97  t_98  t_99  malware  \n",
       "0   76  ...    71   297   135   171   215    35   208    56    71        1  \n",
       "1  172  ...    81   240   117    71   297   135   171   215    35        1  \n",
       "2  240  ...    65   112   123    65   112   123    65   113   112        1  \n",
       "3  172  ...   208   302   208   302   187   208   302   228   302        1  \n",
       "4  117  ...   209   260    40   209   260   141   260   141   260        1  \n",
       "\n",
       "[5 rows x 102 columns]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.read_csv('dynamic_api_call_sequence_per_malware_100_0_306.csv')\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'pandas.core.frame.DataFrame'>\n",
      "RangeIndex: 43876 entries, 0 to 43875\n",
      "Columns: 102 entries, hash to malware\n",
      "dtypes: int64(101), object(1)\n",
      "memory usage: 34.1+ MB\n"
     ]
    }
   ],
   "source": [
    "df.info()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(43876, 100)\n",
      "(43876,)\n"
     ]
    }
   ],
   "source": [
    "X = df.drop(['hash', 'malware'], axis = 1).values.astype(int)\n",
    "y = df['malware'].values.astype(int)\n",
    "print(X.shape)\n",
    "print(y.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "306\n"
     ]
    }
   ],
   "source": [
    "print(X.min())\n",
    "print(X.max())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_imbalance(dataset):\n",
    "    count = sorted(Counter(dataset).items())\n",
    "    print(count)\n",
    "    print(count[1][1] / count[0][1])\n",
    "    return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(0, 1079), (1, 42797)]\n",
      "39.66357738646895\n"
     ]
    }
   ],
   "source": [
    "check_imbalance(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(0, 1079), (1, 1079)]\n",
      "1.0\n"
     ]
    }
   ],
   "source": [
    "random_undersampler = RandomUnderSampler(random_state = SEED)\n",
    "\n",
    "X, y = random_undersampler.fit_resample(X, y)\n",
    "\n",
    "check_imbalance(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, X_test, y_train, y_test = train_test_split(X, y, test_size = 0.3, random_state = SEED)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(0, 738), (1, 772)]\n",
      "1.046070460704607\n",
      "[(0, 341), (1, 307)]\n",
      "0.9002932551319648\n"
     ]
    }
   ],
   "source": [
    "check_imbalance(y_train)\n",
    "check_imbalance(y_test)\n",
    "\n",
    "del df, X, y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Architecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model_2_network(\n",
      "  (dgcnn_1): DGCNN_network()\n",
      "  (dropout): Dropout(p=0.4, inplace=False)\n",
      "  (dgcnn_2): DGCNN_network()\n",
      "  (fc): Linear(in_features=27016, out_features=1, bias=True)\n",
      ")\n",
      "\n",
      "Parameters: 35324\n"
     ]
    }
   ],
   "source": [
    "def norn_adj(X, input_dim_1):\n",
    "    \n",
    "    A = torch.zeros((X.size(0), input_dim_1, input_dim_1), dtype = torch.float).cuda()\n",
    "        \n",
    "    A_view = A.view(A.size(0), -1)\n",
    "    x_size = X.size(-1)\n",
    "    indices = X.narrow(-1, 0, x_size - 1) * A.stride(1) * A.stride(2) + X.narrow(-1, 1, x_size - 1) * A.stride(2)\n",
    "    A_view.scatter_(1, indices, 1)\n",
    "        \n",
    "    A_hat = A + torch.eye(input_dim_1, dtype = torch.float).cuda()\n",
    "    D_hat = A_hat.sum(dim = 1).pow(-1.0).diag_embed()\n",
    "    \n",
    "    return A_hat, D_hat\n",
    "\n",
    "def to_one_hot(X, input_dim_1):\n",
    "    \n",
    "    X = F.one_hot(X, num_classes = input_dim_1).float()    \n",
    "    X = X.permute(0, 2, 1)\n",
    "    \n",
    "    return X\n",
    "\n",
    "class DGCNN_network(nn.Module):\n",
    "    \n",
    "    def __init__(self, weight_dim_1, weight_dim_2):\n",
    "\n",
    "        super(DGCNN_network, self).__init__()\n",
    "        self.weight_dim_1 = weight_dim_1\n",
    "        self.weight_dim_2 = weight_dim_2        \n",
    "        self.weights = nn.Parameter(torch.rand((self.weight_dim_1, weight_dim_2), dtype = torch.float, requires_grad = True))\n",
    "        \n",
    "    def forward(self, A_hat, D_hat, X):\n",
    "        return D_hat.matmul(A_hat).matmul(X).matmul(self.weights)\n",
    "\n",
    "class Model_2_network(nn.Module):\n",
    "    \n",
    "    def __init__(self, input_dim_1, input_dim_2, weight_dim_2, weight_dim_4, dropout_rate):\n",
    "    \n",
    "        super(Model_2_network, self).__init__()\n",
    "        \n",
    "        self.input_dim_1 = input_dim_1\n",
    "        self.input_dim_2 = input_dim_2\n",
    "        self.weight_dim_1 = input_dim_2\n",
    "        self.weight_dim_2 = weight_dim_2\n",
    "        self.weight_dim_3 = weight_dim_2\n",
    "        self.weight_dim_4 = weight_dim_4\n",
    "        self.dropout_rate = dropout_rate\n",
    "        \n",
    "        self.dgcnn_1 = DGCNN_network(self.weight_dim_1, self.weight_dim_2)\n",
    "        self.dropout = nn.Dropout(p = self.dropout_rate)\n",
    "        self.dgcnn_2 = DGCNN_network(self.weight_dim_3, self.weight_dim_4)\n",
    "        self.fc = nn.Linear(self.input_dim_1 * (self.weight_dim_2 + self.weight_dim_4), 1)\n",
    "        \n",
    "    def forward(self, X):\n",
    "        \n",
    "        A_hat, D_hat = norn_adj(X, input_dim_1)\n",
    "        X = to_one_hot(X, input_dim_1)        \n",
    "\n",
    "        H_1 = self.dgcnn_1(A_hat, D_hat, X)\n",
    "        H_1 = self.dropout(H_1)\n",
    "        H_1 = torch.relu(H_1)\n",
    "        H_2 = self.dgcnn_2(A_hat, D_hat, H_1)\n",
    "        H_2 = self.dropout(H_2)\n",
    "        H_2 = torch.relu(H_2)\n",
    "        H_2 = torch.cat([H_1, H_2], 2)\n",
    "        H_2 = H_2.view(H_2.size(0), -1)\n",
    "        H_2 = self.fc(H_2)\n",
    "                \n",
    "        return H_2.squeeze()\n",
    "    \n",
    "model = Model_2_network(\n",
    "    input_dim_1 = 307,\n",
    "    input_dim_2 = 100,\n",
    "    weight_dim_2 = 71,    \n",
    "    weight_dim_4 = 17,\n",
    "    dropout_rate = 0.4\n",
    ")\n",
    "\n",
    "print(model)\n",
    "print(f'\\nParameters: {np.sum([param.numel() for param in model.parameters()])}')\n",
    "del model    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hyperparameter Optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "start_training = time.time()\n",
    "\n",
    "torch.manual_seed(SEED)\n",
    "\n",
    "input_dim_1 = 307\n",
    "input_dim_2 = 100\n",
    "\n",
    "net = NeuralNetBinaryClassifier(\n",
    "    Model_2_network,\n",
    "    module__input_dim_1 = input_dim_1,\n",
    "    module__input_dim_2 = input_dim_2,\n",
    "    module__weight_dim_2 = 17,\n",
    "    module__weight_dim_4 = 17,\n",
    "    module__dropout_rate = 0.4,\n",
    "    batch_size = 32,\n",
    "    max_epochs = 10,\n",
    "    train_split = None,\n",
    "    optimizer = torch.optim.Adam,\n",
    "    iterator_train__shuffle = True,\n",
    "    device = 'cuda'\n",
    ")\n",
    "\n",
    "pipe = Pipeline([\n",
    "    ('net', net)\n",
    "])\n",
    "\n",
    "# LSTM(10) => 12771 => DGCNN(17, 17) => 12428\n",
    "# LSTM(40) => 55881 => DGCNN(71, 71) => 55736\n",
    "# LSTM(70) => 106191 => DGCNN(126, 126) => 105841\n",
    "# LSTM(100) => 163701 => DGCNN(182, 182) => 163073\n",
    "\n",
    "params = {\n",
    "    'net__module__weight_dim_2' : [17, 71, 126, 182],\n",
    "    'net__module__weight_dim_4' : [17, 71, 126, 182],\n",
    "    'net__module__dropout_rate' : [0.4, 0.5, 0.6],\n",
    "    'net__batch_size' : [32, 64, 128],\n",
    "    'net__max_epochs' : [10, 20, 30]\n",
    "}\n",
    "\n",
    "gs = GridSearchCV(\n",
    "    pipe,\n",
    "    params,\n",
    "    refit = False,\n",
    "    cv = StratifiedKFold(n_splits = 5, random_state = SEED, shuffle = True),\n",
    "    scoring = lambda net_gs, X_gs, y_gs : roc_auc_score(y_gs, net_gs.predict_proba(X_gs)),\n",
    "    verbose = 2\n",
    ")\n",
    "\n",
    "gs.fit(X_train, y_train.astype(np.float))\n",
    "\n",
    "print(gs.best_score_, gs.best_params_)\n",
    "\n",
    "end_training = (time.time() - start_training) / 60\n",
    "\n",
    "print(f'\\nDone! Time: {end_training:.2f} min')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Best Configuration:\n",
    "\n",
    "0.9745231287489504 {'net__batch_size': 32, 'net__max_epochs': 30, 'net__module__dropout_rate': 0.6, 'net__module__weight_dim_2': 71, 'net__module__weight_dim_4': 17}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  epoch    train_loss     dur\n",
      "-------  ------------  ------\n",
      "      1       \u001b[36m19.7568\u001b[0m  0.4675\n",
      "      2        \u001b[36m2.7189\u001b[0m  0.2324\n",
      "      3        \u001b[36m0.7381\u001b[0m  0.2291\n",
      "      4        \u001b[36m0.2684\u001b[0m  0.2375\n",
      "      5        \u001b[36m0.2015\u001b[0m  0.2633\n",
      "      6        \u001b[36m0.1484\u001b[0m  0.2399\n",
      "      7        \u001b[36m0.1281\u001b[0m  0.2342\n",
      "      8        \u001b[36m0.1028\u001b[0m  0.2445\n",
      "      9        \u001b[36m0.0969\u001b[0m  0.2410\n",
      "     10        \u001b[36m0.0780\u001b[0m  0.2260\n",
      "     11        0.0802  0.2379\n",
      "     12        0.0786  0.2367\n",
      "     13        \u001b[36m0.0661\u001b[0m  0.2307\n",
      "     14        \u001b[36m0.0615\u001b[0m  0.2280\n",
      "     15        \u001b[36m0.0526\u001b[0m  0.2203\n",
      "     16        \u001b[36m0.0523\u001b[0m  0.2362\n",
      "     17        \u001b[36m0.0519\u001b[0m  0.2397\n",
      "     18        \u001b[36m0.0512\u001b[0m  0.2349\n",
      "     19        \u001b[36m0.0372\u001b[0m  0.2271\n",
      "     20        0.0402  0.2285\n",
      "     21        0.0373  0.2293\n",
      "     22        \u001b[36m0.0327\u001b[0m  0.2260\n",
      "     23        0.0339  0.2170\n",
      "     24        0.0398  0.2225\n",
      "     25        \u001b[36m0.0283\u001b[0m  0.2419\n",
      "     26        0.0327  0.2449\n",
      "     27        0.0351  0.2343\n",
      "     28        0.0300  0.2498\n",
      "     29        0.0368  0.2276\n",
      "     30        0.0319  0.2451\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Pipeline(memory=None,\n",
       "         steps=[('net',\n",
       "                 <class 'skorch.classifier.NeuralNetBinaryClassifier'>[initialized](\n",
       "  module_=Model_2_network(\n",
       "    (dgcnn_1): DGCNN_network()\n",
       "    (dropout): Dropout(p=0.6, inplace=False)\n",
       "    (dgcnn_2): DGCNN_network()\n",
       "    (fc): Linear(in_features=27016, out_features=1, bias=True)\n",
       "  ),\n",
       "))],\n",
       "         verbose=False)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(SEED)\n",
    "\n",
    "# weight_dim_2 = gs.best_params_['net__module__weight_dim_2']\n",
    "# weight_dim_4 = gs.best_params_['net__module__weight_dim_4']\n",
    "# dropout_rate = gs.best_params_['net__module__dropout_rate']\n",
    "# batch_size = gs.best_params_['net__batch_size']\n",
    "# max_epochs = gs.best_params_['net__max_epochs']\n",
    "\n",
    "input_dim_1 = 307\n",
    "input_dim_2 = 100\n",
    "weight_dim_2 = 71\n",
    "weight_dim_4 = 17\n",
    "dropout_rate = 0.6\n",
    "batch_size = 32\n",
    "max_epochs = 30\n",
    "\n",
    "net = NeuralNetBinaryClassifier(\n",
    "    Model_2_network,\n",
    "    module__input_dim_1 = input_dim_1,\n",
    "    module__input_dim_2 = input_dim_2,\n",
    "    module__weight_dim_2 = weight_dim_2,\n",
    "    module__weight_dim_4 = weight_dim_4,\n",
    "    module__dropout_rate = dropout_rate,\n",
    "    batch_size = batch_size,\n",
    "    max_epochs = max_epochs,\n",
    "    train_split = None,\n",
    "    optimizer = torch.optim.Adam,\n",
    "    iterator_train__shuffle = True,\n",
    "    device = 'cuda'\n",
    ")\n",
    "\n",
    "pipe = Pipeline([\n",
    "    ('net', net)\n",
    "])\n",
    "\n",
    "pipe.fit(X_train, y_train.astype(np.float))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_evaluate(y, pred):\n",
    "    print('Confusion matrix\\n[TN FP]\\n[FN TP]')\n",
    "    print(confusion_matrix(y >= 0.5, pred >= 0.5))\n",
    "    print(f'Accuracy: {accuracy_score(y >= 0.5, pred >= 0.5):.4f}')\n",
    "    print(f'Precision: {precision_score(y >= 0.5, pred >= 0.5):.4f}')\n",
    "    print(f'Recall: {recall_score(y >= 0.5, pred >= 0.5):.4f}')\n",
    "    print(f'F1-Score: {f1_score(y >= 0.5, pred >= 0.5):.4f}')\n",
    "    print(f'ROC AUC: {roc_auc_score(y, pred):.4f}')\n",
    "    return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Confusion matrix\n",
      "[TN FP]\n",
      "[FN TP]\n",
      "[[  0 341]\n",
      " [  0 307]]\n",
      "Accuracy: 0.4738\n",
      "Precision: 0.4738\n",
      "Recall: 1.0000\n",
      "F1-Score: 0.6429\n",
      "ROC AUC: 0.5000\n"
     ]
    }
   ],
   "source": [
    "model_evaluate(y_test, np.ones(len(y_test)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Confusion matrix\n",
      "[TN FP]\n",
      "[FN TP]\n",
      "[[317  24]\n",
      " [ 25 282]]\n",
      "Accuracy: 0.9244\n",
      "Precision: 0.9216\n",
      "Recall: 0.9186\n",
      "F1-Score: 0.9201\n",
      "ROC AUC: 0.9729\n"
     ]
    }
   ],
   "source": [
    "model_evaluate(y_test, pipe.predict_proba(X_test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "IPython.notebook.save_notebook()\n"
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%javascript\n",
    "IPython.notebook.save_notebook()"
   ]
  }
 ],
 "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}