blog_notebooks/mortgage_deep_learning/mortgage_e2e_deep_learning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Mortgage Workflow with Deep Learning\n",
    "**Workflow modified with PyTorch DNN by Chris Green**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset\n",
    "\n",
    "The dataset used with this workflow is derived from [Fannie Mae’s Single-Family Loan Performance Data](http://www.fanniemae.com/portal/funding-the-market/data/loan-performance-data.html) with all rights reserved by Fannie Mae. This processed dataset is redistributed with permission and consent from Fannie Mae.\n",
    "\n",
    "To acquire this dataset, please visit [RAPIDS Datasets Homepage](https://rapidsai.github.io/demos/datasets/mortgage-data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Goal\n",
    "The aim of this notebook is to **combine RAPIDS GPU data processing** with a [PyTorch](https://pytorch.org) **deep neural network** (DNN) to train 12-month mortgage loan delinquency prediction model in the same *vein as the XGBoost end-to-end example*.\n",
    "\n",
    "**Note: this notebook is here for archival purposes and is not intended to illustrate best practices or currently working code with the latest RAPIDS Version.  It has been tested to work with RAPIDS 0.7**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ETL: What's new?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The ETL pipeline below looks very much like the original E2E example with two notable exceptions:\n",
    "\n",
    "* Train/Validation/Test Data Split\n",
    "* Feature Discretization & Hashing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Data Split\n",
    "When training a model we want to ensure its generalizability, i.e. that it perorms well on new data. As we train the DNN below we will want to make sure it does not begin to overfit and once trained that we have a good measure of expected performance (if it was to be used in the real-world). To accomplish these ends we break the time series dataset into three non-overlapping fixed time intervals. The earliest will be used as training data, the second as validation data to track the performance of the model during training, and the last as our final test set."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Feature Discretization & Hashing\n",
    "The original ETL pipeline results in both discrete and continuous features. In order to easily employ a DNN we need to transform\n",
    "the continuous features to discrete. This is accomplished by computing a discrete histogram (`Series.quantile`)\n",
    "of these features and then assigning each value to its bin id (`Series.digitize`).\n",
    "\n",
    "Once we have discrete features we want to hash\n",
    "them into a range whose values will act as indices to an embedding table providing\n",
    "the inputs to the network. We use the `Series.hash_encode` method to\n",
    "accomplish this [feature hashing](https://en.wikipedia.org/wiki/Feature_hashing)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PyTorch Deep Neural Network\n",
    "\n",
    "### Model\n",
    "The model constructed below starts with an initial embedding layer ([`torch.nn.EmbeddingBag`](https://pytorch.org/docs/stable/nn.html#embeddingbag)) that takes the indices from the ETL pipeline, looks up the embeddings in the hash table and takes their mean. This vector then passes to a [multilayer perceptron](https://en.wikipedia.org/wiki/Multilayer_perceptron) which finally outputs a single score.\n",
    "\n",
    "Many of the model architecture parameters can be configured by the user such as embedding dimension, number and size of hidden layers, and activation functions.\n",
    "\n",
    "### Training\n",
    "To cut down on boilerplate code and realize the benefits of [early stopping](https://en.wikipedia.org/wiki/Early_stopping)\n",
    "we'll use the [`ignite`](https://pytorch.org/ignite/) library.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Requirements\n",
    "Beyond the dependencies that come installed in the standard \n",
    "[RAPIDS docker containers](https://hub.docker.com/r/rapidsai/rapidsai) we'll also\n",
    "need the following `pip` dependencies installed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting torch\n",
      "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/ac/23/a4b5c189dd624411ec84613b717594a00480282b949e3448d189c4aa4e47/torch-1.1.0-cp37-cp37m-manylinux1_x86_64.whl (676.9MB)\n",
      "\u001b[K     |██████████████████████          | 464.0MB 25.7MB/s eta 0:00:09"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub data rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_data_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_data_rate_limit=1000000.0 (bytes/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[K     |████████████████████████████████| 676.9MB 102kB/s \n",
      "\u001b[?25hCollecting pytorch-ignite\n",
      "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/98/7b/1da69e5fdcb70e8f40ff3955516550207d5f5c81b428a5056510e72c60c5/pytorch_ignite-0.2.0-py2.py3-none-any.whl (73kB)\n",
      "\u001b[K     |████████████████████████████████| 81kB 21.1MB/s \n",
      "\u001b[?25hRequirement already satisfied: numpy in /conda/envs/rapids/lib/python3.7/site-packages (from torch) (1.16.2)\n",
      "Installing collected packages: torch, pytorch-ignite\n",
      "Successfully installed pytorch-ignite-0.2.0 torch-1.1.0\n"
     ]
    }
   ],
   "source": [
    "!pip install torch pytorch-ignite"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CODE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict, OrderedDict\n",
    "import datetime as dt\n",
    "import glob\n",
    "import os\n",
    "import re\n",
    "import subprocess\n",
    "import tempfile\n",
    "import time\n",
    "\n",
    "import cudf\n",
    "import dask\n",
    "from dask.delayed import delayed\n",
    "from dask.distributed import as_completed, Client, wait\n",
    "from dask_cuda import LocalCUDACluster\n",
    "from ignite.engine import create_supervised_evaluator, create_supervised_trainer, Events\n",
    "from ignite.handlers import EarlyStopping as IgniteEarlyStopping\n",
    "from ignite.metrics import Loss, Metric\n",
    "import numpy as np\n",
    "import pyarrow.parquet as pq\n",
    "from sklearn.metrics import auc, precision_recall_curve\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as torch_optim\n",
    "from torch.utils import data as torch_data\n",
    "\n",
    "# CUDF_VERSION = tuple(map(int, cudf.__version__.split(\".\")[:3]))\n",
    "# assert CUDF_VERSION >= (0, 6, 0), \"cudf version must be at least 0.6.0! Found {}!\".format(CUDF_VERSION)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### ETL - Raw Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# to download data for this notebook, visit https://rapidsai.github.io/demos/datasets/mortgage-data and update the following paths accordingly\n",
    "acq_data_path = \"/data/mortgage/acquisition\"\n",
    "perf_data_path = \"/data/mortgage/perf_clean_full_split/\"\n",
    "col_names_path = \"/data/mortgage/names.csv\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### ETL - Data Splits\n",
    "The loan range below is used to select only those loans acquired by FNMA during the configured period. Note that the validation range's\n",
    "earliest date is > 12 months after the last training date because the training target, 12 month delinquency, uses information about one year in the future. To ensure generalizability we must validate and test on data beyond those dates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loans are divided on a per quarter basis\n",
    "loan_range = ['2000Q1', '2000Q4']\n",
    "\n",
    "# Loan performance data are divided on a per month basis\n",
    "train_range = ['200001', '200012']\n",
    "validation_range = ['200201', '200201']\n",
    "test_range = ['200202', '200202']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### ETL - Discretization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_quantiles = 20  # Used for computing histograms of continuous features\n",
    "num_features = 2 ** 18  # When hashing features range will be [0, num_features)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Training - Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_size = 16\n",
    "hidden_dims = [256, 256, 256]\n",
    "\n",
    "device = \"cuda\"\n",
    "dropout = None  # Can add dropout probability in [0, 1] here\n",
    "activation = nn.ReLU()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Training - Optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epoch_size = 10000000\n",
    "\n",
    "train_batch_size = 512\n",
    "validation_batch_size = 2048\n",
    "\n",
    "log_interval = 1000\n",
    "\n",
    "learning_rate = 0.01\n",
    "patience = 4\n",
    "lr_multiplier = 0.5\n",
    "max_epochs = 3  # Increase this for a more realistic training run "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Start Dask CUDA Cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cmd = \"hostname --all-ip-addresses\"\n",
    "process = subprocess.Popen(cmd.split(), stdout=subprocess.PIPE)\n",
    "output, error = process.communicate()\n",
    "IPADDR = str(output.decode()).split()[0]\n",
    "\n",
    "cluster = LocalCUDACluster(ip=IPADDR, scheduler_port=8786, diagnostics_port=8787)\n",
    "client = Client(cluster)\n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_dask_task(func, **kwargs):\n",
    "    task = func(**kwargs)\n",
    "    return task"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Date Range Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_year_quarter(date_str, regex):\n",
    "    matches = regex.findall(date_str)\n",
    "    if not matches:\n",
    "        raise Exception(\"'{}' does not match '{}' pattern!\".format(date_str, regex.pattern))\n",
    "    return int(matches[0][0]), int(matches[0][1])\n",
    "\n",
    "\n",
    "def check_date_range(rng, regex):\n",
    "    dates = [parse_year_quarter(d, regex) for d in rng]\n",
    "    assert dates[0][0] <= dates[1][0], \"First year > second year\"\n",
    "    if dates[0][0] == dates[1][0]:\n",
    "        assert dates[0][1] <= dates[1][1]\n",
    "    return tuple(dates)\n",
    "\n",
    "MONTH_RE = re.compile('([0-9]{4})(0[1-9]|1[0-2])')\n",
    "QUARTER_RE = re.compile('([0-9]{4})Q([1234])')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_date, end_date = check_date_range(loan_range, QUARTER_RE)\n",
    "train_dates = check_date_range(train_range, MONTH_RE)\n",
    "validation_dates = check_date_range(validation_range, MONTH_RE)\n",
    "test_dates = check_date_range(test_range, MONTH_RE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Using data from loans acquired by FNMA from {} to {}\".format(loan_range[0], loan_range[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RMM Pool Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def initialize_rmm_pool():\n",
    "    from librmm_cffi import librmm_config as rmm_cfg\n",
    "\n",
    "    rmm_cfg.use_pool_allocator = True\n",
    "    #rmm_cfg.initial_pool_size = 2<<30 # set to 2GiB. Default is 1/2 total GPU memory\n",
    "    import cudf\n",
    "    return cudf.rmm.initialize()\n",
    "\n",
    "def initialize_rmm_no_pool():\n",
    "    from librmm_cffi import librmm_config as rmm_cfg\n",
    "    \n",
    "    rmm_cfg.use_pool_allocator = False\n",
    "    import cudf\n",
    "    return cudf.rmm.initialize()\n",
    "\n",
    "def finalize_rmm():\n",
    "    import cudf\n",
    "    return cudf.rmm.finalize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.run(initialize_rmm_pool)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Raw Mortgage Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PERFORMANCE_COLS = OrderedDict([\n",
    "    (\"loan_id\", \"int64\"),\n",
    "    (\"monthly_reporting_period\", \"date\"),\n",
    "    (\"servicer\", \"category\"),\n",
    "    (\"interest_rate\", \"float64\"),\n",
    "    (\"current_actual_upb\", \"float64\"),\n",
    "    (\"loan_age\", \"float64\"),\n",
    "    (\"remaining_months_to_legal_maturity\", \"float64\"),\n",
    "    (\"adj_remaining_months_to_maturity\", \"float64\"),\n",
    "    (\"maturity_date\", \"date\"),\n",
    "    (\"msa\", \"float64\"),\n",
    "    (\"current_loan_delinquency_status\", \"int32\"),\n",
    "    (\"mod_flag\", \"category\"),\n",
    "    (\"zero_balance_code\", \"category\"),\n",
    "    (\"zero_balance_effective_date\", \"date\"),\n",
    "    (\"last_paid_installment_date\", \"date\"),\n",
    "    (\"foreclosed_after\", \"date\"),\n",
    "    (\"disposition_date\", \"date\"),\n",
    "    (\"foreclosure_costs\", \"float64\"),\n",
    "    (\"prop_preservation_and_repair_costs\", \"float64\"),\n",
    "    (\"asset_recovery_costs\", \"float64\"),\n",
    "    (\"misc_holding_expenses\", \"float64\"),\n",
    "    (\"holding_taxes\", \"float64\"),\n",
    "    (\"net_sale_proceeds\", \"float64\"),\n",
    "    (\"credit_enhancement_proceeds\", \"float64\"),\n",
    "    (\"repurchase_make_whole_proceeds\", \"float64\"),\n",
    "    (\"other_foreclosure_proceeds\", \"float64\"),\n",
    "    (\"non_interest_bearing_upb\", \"float64\"),\n",
    "    (\"principal_forgiveness_upb\", \"float64\"),\n",
    "    (\"repurchase_make_whole_proceeds_flag\", \"category\"),\n",
    "    (\"foreclosure_principal_write_off_amount\", \"float64\"),\n",
    "    (\"servicing_activity_indicator\", \"category\")\n",
    "])\n",
    "\n",
    "ACQUISITION_COLS = OrderedDict([\n",
    "    (\"loan_id\", \"int64\"),\n",
    "    (\"orig_channel\", \"category\"),\n",
    "    (\"seller_name\", \"category\"),\n",
    "    (\"orig_interest_rate\", \"float64\"),\n",
    "    (\"orig_upb\", \"int64\"),\n",
    "    (\"orig_loan_term\", \"int64\"),\n",
    "    (\"orig_date\", \"date\"),\n",
    "    (\"first_pay_date\", \"date\"),\n",
    "    (\"orig_ltv\", \"float64\"),\n",
    "    (\"orig_cltv\", \"float64\"),\n",
    "    (\"num_borrowers\", \"float64\"),\n",
    "    (\"dti\", \"float64\"),\n",
    "    (\"borrower_credit_score\", \"float64\"),\n",
    "    (\"first_home_buyer\", \"category\"),\n",
    "    (\"loan_purpose\", \"category\"),\n",
    "    (\"property_type\", \"category\"),\n",
    "    (\"num_units\", \"int64\"),\n",
    "    (\"occupancy_status\", \"category\"),\n",
    "    (\"property_state\", \"category\"),\n",
    "    (\"zip\", \"int64\"),\n",
    "    (\"mortgage_insurance_percent\", \"float64\"),\n",
    "    (\"product_type\", \"category\"),\n",
    "    (\"coborrow_credit_score\", \"float64\"),\n",
    "    (\"mortgage_insurance_type\", \"float64\"),\n",
    "    (\"relocation_mortgage_indicator\", \"category\")\n",
    "])\n",
    "    \n",
    "NAMES_COLS = OrderedDict([\n",
    "    (\"seller_name\", \"category\"),\n",
    "    (\"new\", \"category\"),\n",
    "])\n",
    "\n",
    "def gpu_load_performance_csv(performance_path, drop_cols=[], skip_rows=0):\n",
    "    \"\"\" Loads performance data\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    GPU DataFrame\n",
    "    \"\"\"\n",
    "    df = cudf.read_csv(performance_path, names=PERFORMANCE_COLS.keys(), delimiter='|',\n",
    "                       dtype=list(PERFORMANCE_COLS.values()), skiprows=skip_rows)\n",
    "    for col in drop_cols:\n",
    "        df.drop_column(col)\n",
    "    return df\n",
    "\n",
    "def gpu_load_acquisition_csv(acquisition_path, skip_rows=0):\n",
    "    \"\"\" Loads acquisition data\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    GPU DataFrame\n",
    "    \"\"\"\n",
    "    return cudf.read_csv(acquisition_path, names=ACQUISITION_COLS.keys(), \n",
    "                         delimiter='|', dtype=list(ACQUISITION_COLS.values()), skiprows=skip_rows)\n",
    "\n",
    "def gpu_load_names(col_names_path):\n",
    "    \"\"\" Loads names used for renaming the banks\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    GPU DataFrame\n",
    "    \"\"\"\n",
    "    return cudf.read_csv(col_names_path, names=NAMES_COLS.keys(), delimiter='|', \n",
    "                         dtype=list(NAMES_COLS.values()), skiprows=1)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Feature Engineer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_dt_first_day(ym):\n",
    "    return dt.datetime(ym[0], ym[1], 1)\n",
    "\n",
    "\n",
    "def sample_df(df, start, end):\n",
    "\n",
    "    start_dt = make_dt_first_day(start)\n",
    "    end_dt = make_dt_first_day(end)\n",
    "\n",
    "    query = \"timestamp >= @{} and timestamp <= @{}\"\n",
    "    return df.query(query.format(\"start_dt\", \"end_dt\"))\n",
    "\n",
    "\n",
    "def null_workaround(df, **kwargs):\n",
    "    for column, data_type in df.dtypes.items():\n",
    "        if str(data_type) == \"category\":\n",
    "            df[column] = df[column].astype('int32').fillna(np.dtype(np.int32).type(-1))\n",
    "        if str(data_type) in ['int8', 'int16', 'int32', 'int64', 'float32', 'float64']:\n",
    "            df[column] = df[column].fillna(np.dtype(data_type).type(-1)).astype(data_type)\n",
    "    return df\n",
    "\n",
    "\n",
    "def create_ever_features(df):\n",
    "    \"\"\"Calculate if loans ever been delinquent for 30/90/180 days\"\"\"\n",
    "\n",
    "    everdf = df[['loan_id', 'current_loan_delinquency_status']]\n",
    "    everdf = everdf.groupby('loan_id', method='hash', as_index=False).max()\n",
    "    del(df)\n",
    "    everdf['ever_30'] = (everdf['current_loan_delinquency_status'] >= 1).astype('int8')\n",
    "    everdf['ever_90'] = (everdf['current_loan_delinquency_status'] >= 3).astype('int8')\n",
    "    everdf['ever_180'] = (everdf['current_loan_delinquency_status'] >= 6).astype('int8')\n",
    "    everdf.drop_column('current_loan_delinquency_status')\n",
    "    return everdf\n",
    "\n",
    "\n",
    "def create_delinq_features(df, **kwargs):\n",
    "    \"\"\"Find minimum dates when loans are delinquent for 30/90/180 days\"\"\"\n",
    "    delinq_df = df[['loan_id', 'monthly_reporting_period', 'current_loan_delinquency_status']]\n",
    "    del(df)\n",
    "    delinq_30 = delinq_df.query('current_loan_delinquency_status >= 1')[['loan_id', 'monthly_reporting_period']].groupby('loan_id', method='hash', as_index=False).min()\n",
    "    delinq_30['delinquency_30'] = delinq_30['monthly_reporting_period']\n",
    "    delinq_30.drop_column('monthly_reporting_period')\n",
    "    delinq_90 = delinq_df.query('current_loan_delinquency_status >= 3')[['loan_id', 'monthly_reporting_period']].groupby('loan_id', method='hash', as_index=False).min()\n",
    "    delinq_90['delinquency_90'] = delinq_90['monthly_reporting_period']\n",
    "    delinq_90.drop_column('monthly_reporting_period')\n",
    "    delinq_180 = delinq_df.query('current_loan_delinquency_status >= 6')[['loan_id', 'monthly_reporting_period']].groupby('loan_id', method='hash', as_index=False).min()\n",
    "    delinq_180['delinquency_180'] = delinq_180['monthly_reporting_period']\n",
    "    delinq_180.drop_column('monthly_reporting_period')\n",
    "    del(delinq_df)\n",
    "    delinq_merge = delinq_30.merge(delinq_90, how='left', on=['loan_id'], type='hash')\n",
    "    delinq_merge['delinquency_90'] = delinq_merge['delinquency_90'].fillna(np.dtype('datetime64[ms]').type('1970-01-01').astype('datetime64[ms]'))\n",
    "    delinq_merge = delinq_merge.merge(delinq_180, how='left', on=['loan_id'], type='hash')\n",
    "    delinq_merge['delinquency_180'] = delinq_merge['delinquency_180'].fillna(np.dtype('datetime64[ms]').type('1970-01-01').astype('datetime64[ms]'))\n",
    "    del(delinq_30)\n",
    "    del(delinq_90)\n",
    "    del(delinq_180)\n",
    "    return delinq_merge\n",
    "\n",
    "\n",
    "def join_ever_delinq_features(everdf_tmp, delinq_merge, **kwargs):\n",
    "    \"\"\"\n",
    "    Output cols: loan_id|ever_30|ever_90|ever_180|delinquency_30|delinquency_90|delinquency_180\n",
    "    \"\"\"\n",
    "    everdf = everdf_tmp.merge(delinq_merge, on=['loan_id'], how='left', type='hash')\n",
    "    del(everdf_tmp)\n",
    "    del(delinq_merge)\n",
    "    everdf['delinquency_30'] = everdf['delinquency_30'].fillna(np.dtype('datetime64[ms]').type('1970-01-01').astype('datetime64[ms]'))\n",
    "    everdf['delinquency_90'] = everdf['delinquency_90'].fillna(np.dtype('datetime64[ms]').type('1970-01-01').astype('datetime64[ms]'))\n",
    "    everdf['delinquency_180'] = everdf['delinquency_180'].fillna(np.dtype('datetime64[ms]').type('1970-01-01').astype('datetime64[ms]'))\n",
    "    return everdf\n",
    "\n",
    "\n",
    "def create_joined_df(df, everdf, **kwargs):\n",
    "    \"\"\"Join the full dataframe column subset with the joined delinquency/ever dataframe\n",
    "\n",
    "    Note:\n",
    "      * The column 'delinquency_12' returned is just the current loan delinquency status\n",
    "      * The column 'upb_12' returned is just the current actual unpaid principal balance\n",
    "      * Output cols:\n",
    "        ** loan_id\n",
    "    ** timestamp\n",
    "    ** timestamp_month\n",
    "    ** timestamp_year\n",
    "    ** delinquency_12\n",
    "    ** upb_12\n",
    "    ** ever_30\n",
    "    ** ever_90\n",
    "    ** ever_180\n",
    "    ** delinquency_30\n",
    "    ** delinquency_90\n",
    "    ** delinquency_180\n",
    "    \"\"\"\n",
    "    test = df[['loan_id', 'monthly_reporting_period', 'current_loan_delinquency_status', 'current_actual_upb']]\n",
    "    del(df)\n",
    "    test['timestamp'] = test['monthly_reporting_period']\n",
    "    test.drop_column('monthly_reporting_period')\n",
    "    test['timestamp_month'] = test['timestamp'].dt.month\n",
    "    test['timestamp_year'] = test['timestamp'].dt.year\n",
    "\n",
    "    test['delinquency_12'] = test['current_loan_delinquency_status']\n",
    "    test.drop_column('current_loan_delinquency_status')\n",
    "    test['delinquency_12'] = test['delinquency_12'].fillna(-1)\n",
    "\n",
    "    test['upb_12'] = test['current_actual_upb']\n",
    "    test.drop_column('current_actual_upb')\n",
    "    test['upb_12'] = test['upb_12'].fillna(999999999)\n",
    "    \n",
    "    joined_df = test.merge(everdf, how='left', on=['loan_id'], type='hash')\n",
    "    del(everdf)\n",
    "    del(test)\n",
    "    \n",
    "    joined_df['ever_30'] = joined_df['ever_30'].fillna(-1)\n",
    "    joined_df['ever_90'] = joined_df['ever_90'].fillna(-1)\n",
    "    joined_df['ever_180'] = joined_df['ever_180'].fillna(-1)\n",
    "    joined_df['delinquency_30'] = joined_df['delinquency_30'].fillna(-1)\n",
    "    joined_df['delinquency_90'] = joined_df['delinquency_90'].fillna(-1)\n",
    "    joined_df['delinquency_180'] = joined_df['delinquency_180'].fillna(-1)\n",
    "    \n",
    "    joined_df['timestamp_year'] = joined_df['timestamp_year'].astype('int32')\n",
    "    joined_df['timestamp_month'] = joined_df['timestamp_month'].astype('int32')\n",
    "    \n",
    "    return joined_df\n",
    "\n",
    "\n",
    "def create_12_mon_features(joined_df):\n",
    "    testdfs = []\n",
    "    n_months = 12\n",
    "    for y in range(1, n_months + 1):\n",
    "        tmpdf = joined_df[['loan_id', 'timestamp_year', 'timestamp_month', 'delinquency_12', 'upb_12']]\n",
    "        tmpdf['josh_months'] = tmpdf['timestamp_year'] * 12 + tmpdf['timestamp_month']\n",
    "        tmpdf['josh_mody_n'] = ((tmpdf['josh_months'].astype('float64') - 24000 - y) / 12).floor()\n",
    "        tmpdf = tmpdf.groupby(['loan_id', 'josh_mody_n'], method='hash', as_index=False).agg({'delinquency_12': 'max','upb_12': 'min'})\n",
    "        tmpdf['delinquency_12'] = (tmpdf['max_delinquency_12']>3).astype('int32')\n",
    "        tmpdf['delinquency_12'] +=(tmpdf['min_upb_12']==0).astype('int32')\n",
    "        tmpdf.drop_column('max_delinquency_12')\n",
    "        tmpdf['upb_12'] = tmpdf['min_upb_12']\n",
    "        tmpdf.drop_column('min_upb_12')\n",
    "        tmpdf['timestamp_year'] = (((tmpdf['josh_mody_n'] * n_months) + 24000 + (y - 1)) / 12).floor().astype('int16')\n",
    "        tmpdf['timestamp_month'] = np.int8(y)\n",
    "        tmpdf.drop_column('josh_mody_n')\n",
    "        testdfs.append(tmpdf)\n",
    "        del(tmpdf)\n",
    "    del(joined_df)\n",
    "\n",
    "    return cudf.concat(testdfs)\n",
    "\n",
    "\n",
    "def combine_joined_12_mon(joined_df, testdf, **kwargs):\n",
    "    joined_df.drop_column('delinquency_12')\n",
    "    joined_df.drop_column('upb_12')\n",
    "    joined_df['timestamp_year'] = joined_df['timestamp_year'].astype('int16')\n",
    "    joined_df['timestamp_month'] = joined_df['timestamp_month'].astype('int8')\n",
    "    return joined_df.merge(testdf, how='left', on=['loan_id', 'timestamp_year', 'timestamp_month'], type='hash')\n",
    "\n",
    "def final_performance_delinquency(df, joined_df, **kwargs):\n",
    "    merged = null_workaround(df)\n",
    "    joined_df = null_workaround(joined_df)\n",
    "    merged['timestamp_month'] = merged['monthly_reporting_period'].dt.month\n",
    "    merged['timestamp_month'] = merged['timestamp_month'].astype('int8')\n",
    "    merged['timestamp_year'] = merged['monthly_reporting_period'].dt.year\n",
    "    merged['timestamp_year'] = merged['timestamp_year'].astype('int16')\n",
    "    merged = merged.merge(joined_df, how='left', on=['loan_id', 'timestamp_year', 'timestamp_month'], type='hash')\n",
    "\n",
    "    merged.drop_column('timestamp_year')\n",
    "    merged.drop_column('timestamp_month')\n",
    "    return merged\n",
    "\n",
    "\n",
    "def join_perf_acq_dfs(perf, acq, **kwargs):\n",
    "    perf = null_workaround(perf)\n",
    "    acq = null_workaround(acq)\n",
    "    return perf.merge(acq, how='left', on=['loan_id'], type='hash')\n",
    "\n",
    "\n",
    "def last_mile_cleaning(df):\n",
    "    df['timestamp'] = df['monthly_reporting_period'].astype('datetime64[ms]')\n",
    "    drop_list = [\n",
    "        'loan_id', 'orig_date', 'first_pay_date', 'seller_name', 'monthly_reporting_period', \n",
    "        'last_paid_installment_date', 'maturity_date', 'ever_30', 'ever_90', 'ever_180',\n",
    "        'delinquency_30', 'delinquency_90', 'delinquency_180', 'upb_12',\n",
    "        'zero_balance_effective_date','foreclosed_after', 'disposition_date'\n",
    "    ]\n",
    "    for column in drop_list:\n",
    "        df.drop_column(column)\n",
    "    for col, dtype in df.dtypes.iteritems():\n",
    "        if str(dtype)=='category':\n",
    "            df[col] = df[col].cat.codes\n",
    "        if col != 'timestamp':\n",
    "            if 'float' in str(dtype):\n",
    "                df[col] = df[col].astype('float32')\n",
    "    df['delinquency_12'] = df['delinquency_12'] > 0\n",
    "    df['delinquency_12'] = df['delinquency_12'].fillna(False).astype('int8')\n",
    "\n",
    "    for column in df.columns:\n",
    "        df[column] = df[column].fillna(-1)\n",
    "\n",
    "    return df\n",
    "\n",
    "\n",
    "def load_sample_and_clean_data(year,\n",
    "                               quarter,\n",
    "                               perf_file,\n",
    "                               dates=None,\n",
    "                               **kwargs):\n",
    "    names = gpu_load_names(col_names_path)\n",
    "\n",
    "    acq_df = gpu_load_acquisition_csv(acquisition_path=acq_data_path + \"/Acquisition_\"\n",
    "                                      + str(year) + \"Q\" + str(quarter) + \".txt\")\n",
    "    acq_df = acq_df.merge(names, how='left', on=['seller_name'])\n",
    "    acq_df.drop_column('seller_name')\n",
    "    acq_df['seller_name'] = acq_df['new']\n",
    "    acq_df.drop_column('new')\n",
    "    \n",
    "    perf_df_tmp = gpu_load_performance_csv(perf_file)\n",
    "\n",
    "    df = perf_df_tmp\n",
    "    everdf = create_ever_features(df)\n",
    "    delinq_merge = create_delinq_features(df)\n",
    "    everdf = join_ever_delinq_features(everdf, delinq_merge)\n",
    "    del(delinq_merge)\n",
    "    joined_df = create_joined_df(df, everdf)\n",
    "    testdf = create_12_mon_features(joined_df)\n",
    "    joined_df = combine_joined_12_mon(joined_df, testdf)\n",
    "    del(testdf)\n",
    "    perf_df = final_performance_delinquency(df, joined_df)\n",
    "    del(df, joined_df)\n",
    "    final_df = join_perf_acq_dfs(perf_df, acq_df)\n",
    "    del(perf_df)\n",
    "    del(acq_df)\n",
    "    final_df = last_mile_cleaning(final_df)\n",
    "    if dates is None:\n",
    "        final_df.drop_column(\"timestamp\")\n",
    "        out = {\"train\": final_df.to_arrow(preserve_index=False)}\n",
    "        del(final_df)\n",
    "        return out\n",
    "    else:\n",
    "        output = {}\n",
    "        for k, (start, end) in dates.items():\n",
    "            sampled_df = sample_df(final_df, start, end)\n",
    "            sampled_df.drop_column(\"timestamp\")\n",
    "            output[k] = sampled_df.to_arrow(preserve_index=False)\n",
    "            del(sampled_df)\n",
    "        del(final_df)\n",
    "        return output\n",
    "\n",
    "\n",
    "\n",
    "def process_quarter_gpu(client, year, quarter, perf_file, dates=None):\n",
    "\n",
    "    ml_arrays = run_dask_task(delayed(load_sample_and_clean_data),\n",
    "                              year=year,\n",
    "                              quarter=quarter,\n",
    "                              acq_subdir=acq_data_path,\n",
    "                              perf_file=perf_file,\n",
    "                              dates=dates,\n",
    "                              dask_key_name=\"PROCESS/\" + os.path.basename(perf_file))\n",
    "    return client.compute(ml_arrays,\n",
    "                          optimize_graph=False,\n",
    "                          fifo_timeout=\"0ms\")\n",
    "\n",
    "\n",
    "def load_and_process_data(client,\n",
    "                          start_year,\n",
    "                          start_quarter,\n",
    "                          end_year,\n",
    "                          end_quarter,\n",
    "                          dates=None):\n",
    "    assert start_year <= end_year\n",
    "    if start_year == end_year:\n",
    "        assert start_quarter <= end_quarter\n",
    "    \n",
    "    gpu_dfs = []\n",
    "    year = start_year\n",
    "    quarter = start_quarter\n",
    "    while year <= end_year:\n",
    "        if year == end_year and quarter > end_quarter:\n",
    "            break\n",
    "        perf_file_regex = os.path.join(perf_data_path,\n",
    "                                       'Performance_{}Q{}*'.format(year, quarter))\n",
    "        for perf_file in glob.glob(perf_file_regex):\n",
    "            gpu_dfs.append(process_quarter_gpu(client,\n",
    "                                               year,\n",
    "                                               quarter,\n",
    "                                               perf_file,\n",
    "                                               dates=dates))\n",
    "        quarter += 1\n",
    "        if quarter == 5:\n",
    "            year += 1\n",
    "            quarter = 1\n",
    "    return gpu_dfs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dates = {\n",
    "    \"train\": train_dates,\n",
    "    \"validation\": validation_dates,\n",
    "    \"test\": test_dates\n",
    "}\n",
    "\n",
    "tables_fut = load_and_process_data(client, start_date[0], start_date[1], \n",
    "                                   end_date[0], end_date[1], dates=dates)\n",
    "for future in as_completed(tables_fut):\n",
    "    print(future)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compute Discretization & Hashing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@delayed\n",
    "def get_dtype_columns(table_dict, key=\"train\"):\n",
    "    df = cudf.DataFrame.from_arrow(table_dict[key])\n",
    "    columns = defaultdict(list)\n",
    "    items = df.dtypes.items()\n",
    "    del(df)\n",
    "    for column, dtype in items:\n",
    "        columns[str(dtype)].append(column)\n",
    "    return columns\n",
    "\n",
    "\n",
    "@delayed\n",
    "def select_column(tables_dict, column, key=\"train\"):\n",
    "    return tables_dict[key].column(column).to_pandas()\n",
    "\n",
    "\n",
    "def calculate_quantiles(tables, max_quantiles, client, columns, tables_key=\"train\", unique=True):\n",
    "    assert len(tables) > 0\n",
    "    quantiles = {}\n",
    "    for col in columns:\n",
    "        col_pieces = [select_column(td, col, key=tables_key) for td in tables]\n",
    "        col_pieces = client.compute(col_pieces)\n",
    "        wait(col_pieces)\n",
    "        col_list = client.gather(col_pieces)\n",
    "        assert len(col_list)\n",
    "        col_data = np.concatenate(col_list)\n",
    "        series = cudf.Series.from_pandas(col_data, nan_as_null=False)\n",
    "        assert series.null_count == 0, \"Found {} null values. Should be 0!\".format(series.null_count)\n",
    "        step = 1 / max_quantiles\n",
    "        quantiles[col] = series.quantile(np.arange(0, 1, step), quant_index=False).to_array().astype(np.float32)\n",
    "        if unique:\n",
    "            quantiles[col] = np.unique(quantiles[col])\n",
    "    return quantiles\n",
    "\n",
    "\n",
    "@delayed\n",
    "def discretize(tables_dict, bin_dict, hash_size):\n",
    "    for key, table in tables_dict.items():\n",
    "        df = cudf.DataFrame.from_arrow(table)\n",
    "        for col, dtype in df.dtypes.items():\n",
    "            if col != 'delinquency_12':\n",
    "                if col in bin_dict:\n",
    "                    bins = bin_dict[col]\n",
    "                    df[col] = df[col].astype(np.float32).digitize(bins)\n",
    "                    df[col] = df[col].hash_encode(hash_size, use_name=True)\n",
    "                elif 'float' in str(dtype):\n",
    "                    raise RuntimeError(\n",
    "                        \"Float column '{}' does not have bins for discretization!\".format(col))\n",
    "                else:\n",
    "                    df[col] = df[col].hash_encode(hash_size, use_name=True)\n",
    "        tables_dict[key] = df.to_arrow(preserve_index=False)\n",
    "    return tables_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "col_type_fut = client.compute(get_dtype_columns(tables_fut[0]),\n",
    "                              optimize_graph=False, fifo_timeout=\"0ms\")\n",
    "wait(col_type_fut)\n",
    "col_type_dict = client.gather(col_type_fut)\n",
    "cts_columns = col_type_dict[\"float32\"] + col_type_dict[\"float64\"]\n",
    "print(\"Found {} continuous value columns: {}\".format(len(cts_columns),\n",
    "                                                     \",\".join(cts_columns)))\n",
    "quantiles = calculate_quantiles(tables_fut, max_quantiles, client, cts_columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bin_delayed = []\n",
    "for df in tables_fut:\n",
    "    key = \"DISCRETIZE/\" + os.path.basename(df.key)\n",
    "    task = run_dask_task(discretize,\n",
    "                         tables_dict=df,\n",
    "                         bin_dict=quantiles,\n",
    "                         hash_size=num_features,\n",
    "                         dask_key_name=key)\n",
    "    bin_delayed.append(task)\n",
    "\n",
    "bin_fut = client.compute(bin_delayed, optimize_graph=False, fifo_timeout=\"0ms\")\n",
    "for future in as_completed(bin_fut):\n",
    "    print(future)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Persist Data in Apache Parquet Format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@delayed\n",
    "def save(tables_dict, base_output_dir, fnames):\n",
    "    out_dirs = {}\n",
    "    for k in tables_dict.keys():\n",
    "        out_dirs[k] = os.path.join(base_output_dir, k)\n",
    "        assert os.path.exists(out_dirs[k]), \"Output directory {} does not exist!\".format(out_dirs[k])\n",
    "\n",
    "    for k, table in tables_dict.items():\n",
    "        path = os.path.join(base_output_dir, k, fnames[k])\n",
    "        assert not os.path.exists(path), \"Output path already exists at {}!\".format(path)\n",
    "        pq.write_table(table, path, compression='snappy')\n",
    "\n",
    "        \n",
    "def persist_data(client, tables_fut, out_fname_dates):\n",
    "    out_delayed = []\n",
    "    for td in tables_fut:\n",
    "        perf_file = os.path.basename(td.key)\n",
    "        key = os.path.join(\"PERSIST_DNN\", perf_file)\n",
    "        fnames = {k: perf_file.replace(\".\", \"_\") + \"_\" + v + \".parquet\" for (k, v) in out_fname_dates.items()}\n",
    "        task = run_dask_task(save,\n",
    "                             tables_dict=td,\n",
    "                             base_output_dir=out_dir,\n",
    "                             fnames=fnames,\n",
    "                             dask_key_name=key)\n",
    "        out_delayed.append(task)\n",
    "    out_fut = client.compute(out_delayed, optimize_graph=False, fifo_timeout=\"0ms\")\n",
    "    for future in as_completed(out_fut):\n",
    "        print(future)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out_fname_dates = {}\n",
    "out_dir = tempfile.mkdtemp()\n",
    "for key, ((sy, sm), (ey, em)) in dates.items():\n",
    "    key_dir = os.path.join(out_dir, key)\n",
    "    print(\"Creating directory: {}\".format(key_dir))\n",
    "    os.makedirs(key_dir)\n",
    "    out_fname_dates[key] = \"{}{:02d}_{}{:02d}\".format(sy, sm, ey, em)\n",
    "persist_data(client, bin_fut, out_fname_dates)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Shutdown Dask - ETL Complete"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.close()\n",
    "cluster.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Torch Dataset from Parquet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_tensors_from_parquet(path, target_name='delinquency_12'):\n",
    "    tbl = pq.read_table(path).to_pandas()\n",
    "    target = None\n",
    "    if target_name in tbl:\n",
    "        target = torch.from_numpy(tbl.pop(target_name).values.astype(np.float32))\n",
    "    features = torch.from_numpy(tbl.values.astype(np.long))\n",
    "    tensors = [features]\n",
    "    if target is not None:\n",
    "        tensors.append(target)\n",
    "    return tuple(tensors)\n",
    "\n",
    "\n",
    "class MortgageParquetDataset(torch_data.Dataset):\n",
    "\n",
    "    def __init__(self, root_path, num_samples=None, target_name='delinquency_12',\n",
    "                 shuffle_files=False):\n",
    "        self.parquet_files = glob.glob(os.path.join(root_path, \"*.parquet\"))\n",
    "        if shuffle_files:\n",
    "            self.parquet_files = list(np.random.permutation(self.parquet_files))\n",
    "        self.target_name = target_name\n",
    "        self.metadata = [pq.read_metadata(p) for p in self.parquet_files]\n",
    "        self.cumsum_rows = np.cumsum([m.num_rows for m in self.metadata])\n",
    "\n",
    "        self.times_through = 0\n",
    "        if num_samples is not None:\n",
    "            self.num_samples = min(num_samples, self.cumsum_rows[-1])\n",
    "        else:\n",
    "            self.num_samples = self.cumsum_rows[-1]\n",
    "\n",
    "        self.loaded_tensors = None\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.num_samples\n",
    "\n",
    "    def __getitem__(self, item):\n",
    "        tt = self.times_through\n",
    "        if item == len(self) - 1:\n",
    "            self.times_through += 1\n",
    "        item += tt * len(self)\n",
    "        item %= len(self)\n",
    "\n",
    "        part_idx = np.searchsorted(self.cumsum_rows, item, side='right')\n",
    "\n",
    "        if self.loaded_tensors is None or self.loaded_tensors[0] != part_idx:\n",
    "            tensors = load_tensors_from_parquet(self.parquet_files[part_idx])\n",
    "            self.loaded_tensors = (part_idx, tensors)\n",
    "\n",
    "        i = item if part_idx == 0 else item - self.cumsum_rows[part_idx - 1]\n",
    "        return tuple(tensor[i] for tensor in self.loaded_tensors[1])\n",
    "\n",
    "\n",
    "\n",
    "def load_torch_dataset(root_path, num_samples=None, shuffle_files=False):\n",
    "    return MortgageParquetDataset(root_path, num_samples=num_samples, shuffle_files=shuffle_files)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PyTorch DNN Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _make_hidden_layer(in_dim, out_dim, activation, dropout=None):\n",
    "    if dropout:\n",
    "        return nn.Sequential(nn.Linear(in_dim, out_dim), activation, nn.Dropout(p=dropout))\n",
    "    return nn.Sequential(nn.Linear(in_dim, out_dim), activation)\n",
    "\n",
    "\n",
    "class MortgageNetwork(nn.Module):\n",
    "    \"\"\"Mortgage Delinquency DNN.\"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        num_features,\n",
    "        embedding_size,\n",
    "        hidden_dims,\n",
    "        use_cuda=True,\n",
    "        activation=nn.ReLU(),\n",
    "        dropout=None,\n",
    "        embedding_bag_mode='mean'\n",
    "    ):\n",
    "        super(MortgageNetwork, self).__init__()\n",
    "        self.input_size = num_features\n",
    "        self.embedding_size = embedding_size\n",
    "        if use_cuda and torch.cuda.is_available():\n",
    "            self.device = torch.device(\"cuda\")\n",
    "        else:\n",
    "            self.device = torch.device(\"cpu\")\n",
    "        self.activation = activation\n",
    "        self.dropout = dropout\n",
    "\n",
    "        self.embedding = nn.modules.EmbeddingBag(self.input_size, self.embedding_size,\n",
    "                                                 mode=embedding_bag_mode)\n",
    "\n",
    "        if len(hidden_dims) > 0:\n",
    "            dims = [self.embedding_size] + hidden_dims\n",
    "            hidden_layers = [\n",
    "                _make_hidden_layer(dims[i], dims[i + 1], self.activation, self.dropout)\n",
    "                for i in range(len(dims) - 1)\n",
    "            ]\n",
    "            self.hidden_layers = nn.ModuleList(hidden_layers)\n",
    "            self.hidden_layers.extend([nn.Linear(dims[-1], 1)])\n",
    "        else:\n",
    "            self.hidden_layers = []\n",
    "\n",
    "        self.to(self.device)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"Forward pass.\"\"\"\n",
    "        out = self.embedding(x)\n",
    "        out = self.activation(out)\n",
    "        for layer in self.hidden_layers:\n",
    "            out = layer(out)\n",
    "        return out.squeeze()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Metric Used for Early Stopping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PrAucMetric(Metric):\n",
    "    def __init__(self, ignore_bad_metric=False):\n",
    "        super(PrAucMetric, self).__init__()\n",
    "        self.name = \"PR-AUC\"\n",
    "        self._predictions = []\n",
    "        self._targets = []\n",
    "        self._ignore_bad_metric = ignore_bad_metric\n",
    "\n",
    "    def reset(self):\n",
    "        self._predictions = []\n",
    "        self._targets = []\n",
    "\n",
    "    def update(self, output):\n",
    "        if len(output) == 2:\n",
    "            y_pred, y_target = output\n",
    "        else:\n",
    "            raise Exception(\"Expected output of length 2!\")\n",
    "        self._predictions.append(y_pred)\n",
    "        self._targets.append(y_target)\n",
    "\n",
    "    def curve(self, targets, predictions):\n",
    "        prec, rec, _ = precision_recall_curve(targets, predictions)\n",
    "        return rec, prec, None\n",
    "\n",
    "    def compute(self):\n",
    "        targets = torch.cat(self._targets).cpu()\n",
    "        predictions = torch.cat(self._predictions).cpu()\n",
    "        print(\"Number of targets for {}-Curve: {}\".format(self.name, len(targets)))\n",
    "        start = time.time()\n",
    "        x, y, _ = self.curve(targets, predictions)\n",
    "        if not self._ignore_bad_metric and len(x) == 2:\n",
    "            raise MetricCurveError(\"{}-Curve returned only two points!\".format(self.name))\n",
    "        start = time.time()\n",
    "        output = auc(x, y)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Early Stopping Handler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EarlyStopping(IgniteEarlyStopping):\n",
    "    def __init__(\n",
    "        self, model, optimizer, lr_multiplier=0.5, min_lr=1.0e-7, delta=0.0005, *args, **kwargs\n",
    "    ):\n",
    "        super(EarlyStopping, self).__init__(*args, **kwargs)\n",
    "        self.optimizer = optimizer\n",
    "        self.model = model\n",
    "        self.lr_multiplier = lr_multiplier\n",
    "        self.min_lr = min_lr\n",
    "        tmp_dir = tempfile.mkdtemp()\n",
    "        self._state_path = os.path.join(tmp_dir, \"best_state.pth\")\n",
    "        self.delta = delta\n",
    "\n",
    "    def _state(self):\n",
    "        return {\n",
    "            \"model\": self.model.state_dict(),\n",
    "            \"optimizer\": self.optimizer.state_dict(),\n",
    "        }\n",
    "\n",
    "    def _save_state(self):\n",
    "        print(\"Saving state to {}.\".format(self._state_path))\n",
    "        state = self._state()\n",
    "        torch.save(state, self._state_path)\n",
    "\n",
    "    def _load_state(self, update_lr=True):\n",
    "        print(\"Loading state from {}.\".format(self._state_path))\n",
    "        state = torch.load(self._state_path)\n",
    "        self.model.load_state_dict(state[\"model\"])\n",
    "\n",
    "        new_lr = max(self.optimizer.param_groups[0][\"lr\"] * self.lr_multiplier, self.min_lr)\n",
    "        self.optimizer.load_state_dict(state[\"optimizer\"])\n",
    "        if update_lr:\n",
    "            self.optimizer.param_groups[0][\"lr\"] = new_lr\n",
    "            self._logger.info(\"Updated optimizer: {}\".format(str(self.optimizer)))\n",
    "\n",
    "\n",
    "    def __call__(self, engine):\n",
    "        score = self.score_function(engine)\n",
    "\n",
    "        if self.best_score is None:\n",
    "            self.best_score = score\n",
    "            self._save_state()\n",
    "        elif score < self.best_score + self.delta:\n",
    "            self.counter += 1\n",
    "            print(\"Score did not improve! EarlyStopping: %i / %i\" % (self.counter, self.patience))\n",
    "            self._load_state()\n",
    "            if self.counter >= self.patience:\n",
    "                print(\"EarlyStopping: Stop training\")\n",
    "                self.trainer.terminate()\n",
    "\n",
    "        else:\n",
    "            self.best_score = score\n",
    "            self.counter = 0\n",
    "            self._save_state()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_training(model):\n",
    "    # Data\n",
    "    train_dataset = load_torch_dataset(os.path.join(out_dir, \"train\"), epoch_size, shuffle_files=True)\n",
    "    validation_dataset = load_torch_dataset(os.path.join(out_dir, \"validation\"))\n",
    "    test_dataset = load_torch_dataset(os.path.join(out_dir, \"test\"))\n",
    "    \n",
    "    train_loader = torch_data.DataLoader(train_dataset,\n",
    "                                     batch_size=train_batch_size,\n",
    "                                     num_workers=8)\n",
    "    validation_loader = torch_data.DataLoader(validation_dataset,\n",
    "                                         batch_size=validation_batch_size,\n",
    "                                         num_workers=8)\n",
    "    test_loader = torch_data.DataLoader(test_dataset,\n",
    "                                        batch_size=validation_batch_size,\n",
    "                                        num_workers=8)\n",
    "    # Optimizer\n",
    "    optimizer = torch_optim.Adam(model.parameters(), lr=learning_rate)\n",
    "    \n",
    "    # Loss Function\n",
    "    loss_fn = lambda pred, target: F.binary_cross_entropy_with_logits(pred, target)\n",
    "\n",
    "    trainer = create_supervised_trainer(model=model, optimizer=optimizer, loss_fn=loss_fn, device=device)\n",
    "    evaluator = create_supervised_evaluator(model, metrics={\"pr-auc\": PrAucMetric()}, device=device)\n",
    "\n",
    "    # Early stopping\n",
    "    early_stopping_handler = EarlyStopping(\n",
    "        model=model,\n",
    "        optimizer=optimizer,\n",
    "        lr_multiplier=lr_multiplier,\n",
    "        patience=patience,\n",
    "        score_function=lambda engine: engine.state.metrics[\"pr-auc\"],\n",
    "        trainer=trainer,)\n",
    "    evaluator.add_event_handler(Events.COMPLETED, early_stopping_handler)\n",
    "\n",
    "    # Events\n",
    "    @trainer.on(Events.EPOCH_STARTED)\n",
    "    def timer(engine):\n",
    "        setattr(engine.state, \"epoch_start\", time.time())\n",
    "\n",
    "    num_epoch_batches = len(train_loader)\n",
    "    examples_per_epoch = num_epoch_batches * train_batch_size\n",
    "    @trainer.on(Events.ITERATION_COMPLETED)\n",
    "    def log_training_loss(engine):\n",
    "        iter = (engine.state.iteration - 1) % num_epoch_batches + 1\n",
    "        if iter % log_interval == 0:\n",
    "            epoch_time_elapsed = time.time() - engine.state.epoch_start\n",
    "            examples = engine.state.iteration * train_batch_size\n",
    "            epoch_examples_per_second = (examples - (engine.state.epoch - 1) * examples_per_epoch) / epoch_time_elapsed\n",
    "            print(\n",
    "                \"Epoch[{}] Iteration[{}/{}] Loss: {:.5f} Example/s: {:.3f} (Total examples: {})\".format(\n",
    "                    engine.state.epoch, iter, num_epoch_batches, engine.state.output,\n",
    "                    epoch_examples_per_second, examples))\n",
    "\n",
    "    @trainer.on(Events.EPOCH_COMPLETED)\n",
    "    def log_validation_results(engine):\n",
    "        evaluator.run(validation_loader)\n",
    "        metrics = evaluator.state.metrics\n",
    "        pr_auc = metrics[\"pr-auc\"]\n",
    "        print(\"Validation Results - Epoch: {}\\n\\tPR-AUC: {:.5f}\".format(engine.state.epoch, pr_auc))\n",
    "\n",
    "    @trainer.on(Events.COMPLETED)\n",
    "    def log_test_results(engine):\n",
    "        evaluator.run(test_loader)\n",
    "        metrics = evaluator.state.metrics\n",
    "        pr_auc = metrics[\"pr-auc\"]\n",
    "        print(\"Final Test Results - PR-AUC: {:.5f}\".format(pr_auc))\n",
    "    trainer.run(train_loader, max_epochs=max_epochs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = MortgageNetwork(num_features, embedding_size, hidden_dims,\n",
    "                        dropout=dropout, activation=activation)\n",
    "run_training(model)"
   ]
  }
 ],
 "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}