00_core.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp core"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TabNet functionality\n",
    "\n",
    "> Integrating TabNet with fastai."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from fastai.basics import *\n",
    "from fastai.tabular.all import *\n",
    "from pytorch_tabnet.tab_network import TabNetNoEmbeddings\n",
    "from pytorch_tabnet.utils import create_explain_matrix\n",
    "from scipy.sparse import csc_matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TabNetModel(Module):\n",
    "    \"Attention model for tabular data.\"\n",
    "    def __init__(self, emb_szs, n_cont, out_sz, embed_p=0., y_range=None, \n",
    "                 n_d=8, n_a=8,\n",
    "                 n_steps=3, gamma=1.3, \n",
    "                 n_independent=2, n_shared=2, epsilon=1e-15,\n",
    "                 virtual_batch_size=128, momentum=0.02,\n",
    "                 mask_type=\"sparsemax\"):\n",
    "        self.emb_szs = emb_szs\n",
    "        self.embeds = nn.ModuleList([Embedding(ni, nf) for ni,nf in emb_szs])\n",
    "        self.emb_drop = nn.Dropout(embed_p)\n",
    "        self.bn_cont = nn.BatchNorm1d(n_cont)\n",
    "        n_emb = sum(e.embedding_dim for e in self.embeds)\n",
    "        self.n_emb,self.n_cont,self.y_range = n_emb,n_cont,y_range\n",
    "        self.tab_net = TabNetNoEmbeddings(n_emb + n_cont, out_sz, n_d, n_a, n_steps, \n",
    "                                          gamma, n_independent, n_shared, epsilon, \n",
    "                                          virtual_batch_size, momentum, mask_type)\n",
    "        # embeddings dim reducer for explainability \n",
    "        cat_dims = L(emb[1] for emb in emb_szs)\n",
    "        pre_emb = n_cont + len(emb_szs)\n",
    "        self.post_emb = n_cont + n_emb\n",
    "        cat_idxs = list(range(len(emb_szs)))\n",
    "        self.emb_reducer = create_explain_matrix(pre_emb,\n",
    "                                            cat_dims,\n",
    "                                            cat_idxs,\n",
    "                                            self.post_emb)\n",
    "\n",
    "    def embedder(self, x_cat, x_cont):\n",
    "        if self.n_emb != 0:\n",
    "            x = [e(x_cat[:,i]) for i,e in enumerate(self.embeds)]\n",
    "            x = torch.cat(x, 1)\n",
    "            x = self.emb_drop(x)\n",
    "        if self.n_cont != 0:\n",
    "            x_cont = self.bn_cont(x_cont)\n",
    "            x = torch.cat([x, x_cont], 1) if self.n_emb != 0 else x_cont\n",
    "        return x\n",
    "\n",
    "    def forward(self, x_cat, x_cont):\n",
    "        x = self.embedder(x_cat, x_cont)\n",
    "        y, _ = self.tab_net(x)\n",
    "        if self.y_range is not None:\n",
    "            y = (self.y_range[1]-self.y_range[0]) * torch.sigmoid(y) + self.y_range[0]\n",
    "        return y\n",
    "\n",
    "    def forward_masks(self, x_cat, x_cont):\n",
    "        x = self.embedder(x_cat, x_cont)\n",
    "        return self.tab_net.forward_masks(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def tabnet_explain(model:TabNetModel, dl:TabDataLoader):\n",
    "    \"Get explain values for a set of predictions\"\n",
    "    dec_y = []\n",
    "    model.eval()\n",
    "    for batch_nb, data in enumerate(dl):\n",
    "        with torch.no_grad():\n",
    "            M_explain, masks = model.forward_masks(data[0], data[1])\n",
    "        for key, value in masks.items():\n",
    "            masks[key] = csc_matrix.dot(value.cpu().numpy(), model.emb_reducer)\n",
    "\n",
    "        explain = csc_matrix.dot(M_explain.cpu().numpy(), model.emb_reducer)\n",
    "        if batch_nb == 0:\n",
    "            res_explain = explain\n",
    "            res_masks = masks\n",
    "        else:\n",
    "            res_explain = np.vstack([res_explain, explain])                                     \n",
    "            for key, value in masks.items():\n",
    "                res_masks[key] = np.vstack([res_masks[key], value])\n",
    "    return res_explain, res_masks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def tabnet_feature_importances(model:TabNetModel, dl:TabDataLoader):\n",
    "    model.eval()\n",
    "    feature_importances_ = np.zeros((model.post_emb))\n",
    "    for batch_nb, data in enumerate(dl):\n",
    "        M_explain, masks = model.forward_masks(data[0], data[1])\n",
    "        feature_importances_ += M_explain.sum(dim=0).cpu().detach().numpy()\n",
    "\n",
    "    feature_importances_ = csc_matrix.dot(\n",
    "        feature_importances_, model.emb_reducer\n",
    "    )\n",
    "    return feature_importances_ / np.sum(feature_importances_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def size(x:nn.Module, with_grad: bool=True) -> Int:\n",
    "    return sum(p.numel() for p in x.parameters() if p.requires_grad or not with_grad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Converted 00_core.ipynb.\n",
      "Converted 01_examples.ipynb.\n",
      "Converted index.ipynb.\n"
     ]
    }
   ],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *\n",
    "from nbdev.export import notebook2script\n",
    "notebook2script() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:fast_tabnet]",
   "language": "python",
   "name": "conda-env-fast_tabnet-py"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}