{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":" # <h1 style='background:#F7B2B0; border:0; color:black'><center>THE GATEDTABTRANSFORMER - AN ENHANCED DEEP LEARNING ARCHITECTURE FOR TABULAR MODELING.</center></h1> ","metadata":{}},{"cell_type":"markdown","source":"This is a two part tutorial series containing the implementation of [GatedTabTransformer](https://arxiv.org/pdf/2201.00199.pdf) paper in both FLAX and TensorFlow (TPU)\n\n[Part 1 : GatedTabTransformer in FLAX](https://www.kaggle.com/code/usharengaraju/gatedtabtransformer-flax)\n\n[Part 2 : GatedTabTransformer in TensorFlow + TPU](https://www.kaggle.com/code/usharengaraju/tensorflow-tpu-gatedtabtransformer)\n\n","metadata":{}},{"cell_type":"markdown","source":"### Abstract\n\nOver the last few years the research towards using deep learning for tabular data has been on the rise .The state of the art TabTransformer incorporates an attention mechanism to better track relationships between categorical features and then makes use of a standard MLP to output its final logits.GatedTabTransformer implements linear projections are implemented in the MLP block and the paper also experiments with several activation functions .\n\n### Introduction\n\nTabular data is the most commonly used data type in real world applications . Tree based ensemble methods like LightGBM , XGBoost are the current state of the art approaches for tabular data . Over the last few years , there is increasing interest in the usage of deep learning techniques for tabular data primarily because of the eliminating the need for manual embedding and feature engineering . Some of the neural networks architectures which have comparable performance with Tree based ensemble methods are TabNet , DNF-Net etc.\n\nThere is an increasing usage of attention-based architectures like Transformers which was originally used to  handle NLP tasks to solve tabular data problems . TabTransformer is one such architecture which focuses on using Multi-Head Self Attention blocks to model relationships between the categorical features in tabular data, transforming them into robust contextual embeddings.The transformed categorical features are concatenated with continuous values and then fed through a standard multilayer perceptron which makes TabTransformer significantly outperform other deep learning counterparts like TabNet and MLP . GatedTabTransformer further enhances TabTransformer by replacing the final MLP block with a gated multi-layer perceptron (gMLP) , a simple MLP-based network with spatial gating projections, which aims to be on par with Transformers in terms of performance on sequential data\n\n\n![](https://i.imgur.com/nw9hqQW.jpg)\n\n\n### TabTransformer    \n\nThe TabTransformer model, outperforms the other state-of-the-art deep learning methods for tabular data by at least 1.0% on mean AUROC. It consists of a column embedding layer, a stack of N Transformer layers, and a multilayer perceptron . The inputted tabular features are split in two parts for the categorical and continuous values. Column embedding is performed for each categorical feature .It generates parametric embeddings which are inputted to a stack of Transformer layers. Each Transformer layer consists of a multi-head self-attention layer followed by a position-wise feed-forward layer. After the processing of categorical values , they are concatenated along with the continuous values to form a final feature vector which is inputted to a standard multilayer perceptron\n\n![](https://i.imgur.com/mnv2bLy.png)\n\n\n### gMLP model\n\nThe gMLP model consists of a stack of multiple identically structured blocks ,activation function and linear projections along the channel dimension and the spatial gating unit which captures spatial cross-token interactions. The weights are initialized as near-zero values and the biases as ones at the beginning of training. This structure does not require positional embeddings because relevant information will be captured in the gating units. gMLP has been proposed as an alternative to Transformers for NLP and vision tasks having up to 66% less trainable parameters. GatedTabTransformers replaces the pure MLP block in the TabTransformer with gMLP \n\n![](https://i.imgur.com/SRyVmYY.png)\n\n\n\n### GatedTabTransformer\n\nThe column embeddings are generated from categorical data features and continuous values are passed through a normalization layer. The categorical embeddings are then processed by a Transformer block. \n\nTransformer block represents the encoder part of a Transformer . It has two sub-layers - a multi-head self-attention mechanism, and a simple, position wise fully connected feed-forward network. In the final layer MLP is replaced by gMLP and the architecture is adapted to output classification logits and works best for optimization of cross entropy or binary cross entropy loss \n\n![](https://i.imgur.com/ROdKcQy.png)\n","metadata":{}},{"cell_type":"markdown","source":"Reference : https://flax.readthedocs.io/en/latest/overview.html\n\n**FLAX**\n\nFlax is a high-performance neural network library for JAX that is designed for flexibility: Try new forms of training by forking an example and by modifying the training loop, not by adding features to a framework.\n\nFlax is being developed in close collaboration with the JAX team and comes with everything you need to start your research, including:\n\n**Neural network API** (flax.linen): Dense, Conv, {Batch|Layer|Group} Norm, Attention, Pooling, {LSTM|GRU} Cell, Dropout\n\n**Utilities and patterns**: replicated training, serialization and checkpointing, metrics, prefetching on device\n\n**Educational examples that work out of the box**: MNIST, LSTM seq2seq, Graph Neural Networks, Sequence Tagging\n\n**Fast, tuned large-scale end-to-end examples**: CIFAR10, ResNet on ImageNet, Transformer LM1b\n\n\n![](https://i.imgur.com/PQvbMNo.png)","metadata":{}},{"cell_type":"code","source":"pip install --quiet jax flax chex typing","metadata":{"id":"ZNTos0k-9Vwq","outputId":"45db369d-b37c-4018-d828-e2c4b9553f3f","execution":{"iopub.status.busy":"2022-11-21T18:48:41.194622Z","iopub.execute_input":"2022-11-21T18:48:41.195485Z","iopub.status.idle":"2022-11-21T18:48:57.621804Z","shell.execute_reply.started":"2022-11-21T18:48:41.195387Z","shell.execute_reply":"2022-11-21T18:48:57.620413Z"},"_kg_hide-output":true,"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import jax\nimport flax\nfrom typing import Sequence, Type\nfrom flax import linen as nn\nfrom flax.linen.module import Module\nfrom functools import partial\nfrom typing import Any\nimport jax.numpy as jnp\nfrom chex import Array\nfrom jax import random\nDtype = Any\n__all__ = [\"Attention\", \"SpatialGatingUnit\", \"LayerNorm\",\"Sequential\", \"Residual\", \"PreNorm\", \"Identity\"]\nATTN_MASK_VALUE = -1e10\nDtype = Any\nLayerNorm = partial(nn.LayerNorm)","metadata":{"id":"uCFa8pjjH6j9","execution":{"iopub.status.busy":"2022-11-21T18:48:59.694022Z","iopub.execute_input":"2022-11-21T18:48:59.694630Z","iopub.status.idle":"2022-11-21T18:48:59.854411Z","shell.execute_reply.started":"2022-11-21T18:48:59.694594Z","shell.execute_reply":"2022-11-21T18:48:59.853246Z"},"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":" # <h1 style='background:#F7B2B0; border:0; color:black'><center>Exploration of Data</center></h1> ","metadata":{}},{"cell_type":"code","source":"import pandas as pd\ndf  = pd.read_csv('../input/amex-tfrecords/minidata (1).csv')\ncat_features  = ['B_30', 'B_38', 'D_63', 'D_64', 'D_66', 'D_68', 'D_114', 'D_116', 'D_117', 'D_120', 'D_126']\ncont_features = [x for x in list(df.columns) if x not in cat_features]\ncont_features.remove('target')\nlen(cont_features)","metadata":{"id":"hJIJkLqV2kaf","outputId":"bdf3508c-85a1-44d1-8b6a-0901794ccaaf","execution":{"iopub.status.busy":"2022-11-21T18:48:59.856072Z","iopub.execute_input":"2022-11-21T18:48:59.856923Z","iopub.status.idle":"2022-11-21T18:49:01.398897Z","shell.execute_reply.started":"2022-11-21T18:48:59.856889Z","shell.execute_reply":"2022-11-21T18:49:01.397468Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(cat_features)","metadata":{"id":"9Xftvb91HKyv","outputId":"dd6e91b0-174a-47d8-e5ce-8ffd29d5e9b3","execution":{"iopub.status.busy":"2022-11-21T18:49:01.401813Z","iopub.execute_input":"2022-11-21T18:49:01.402464Z","iopub.status.idle":"2022-11-21T18:49:01.409662Z","shell.execute_reply.started":"2022-11-21T18:49:01.402424Z","shell.execute_reply":"2022-11-21T18:49:01.408590Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cats = []\nfor col in cat_features:\n  cats.append(df[col].unique().shape[0])\ncats","metadata":{"id":"MjnJUqlhIVan","outputId":"c242c5d7-4770-436e-96f3-38463eef925b","execution":{"iopub.status.busy":"2022-11-21T18:49:01.411267Z","iopub.execute_input":"2022-11-21T18:49:01.412009Z","iopub.status.idle":"2022-11-21T18:49:01.439630Z","shell.execute_reply.started":"2022-11-21T18:49:01.411973Z","shell.execute_reply":"2022-11-21T18:49:01.438253Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cont_arr = jnp.array(df[cont_features].to_numpy())\ncat_arr = jnp.array(df[cat_features].to_numpy(),dtype=int)\ny_arr = jnp.array(df['target'].to_numpy())\nprint(cont_arr.shape,cat_arr.shape,y_arr.shape)","metadata":{"id":"Adf6gTPkHyro","outputId":"e9b14fc9-ac40-44c7-d221-8e7a1d782981","execution":{"iopub.status.busy":"2022-11-21T18:49:01.441774Z","iopub.execute_input":"2022-11-21T18:49:01.442252Z","iopub.status.idle":"2022-11-21T18:49:01.539270Z","shell.execute_reply.started":"2022-11-21T18:49:01.442202Z","shell.execute_reply":"2022-11-21T18:49:01.538436Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"\n<img src=\"https://camo.githubusercontent.com/dd842f7b0be57140e68b2ab9cb007992acd131c48284eaf6b1aca758bfea358b/68747470733a2f2f692e696d6775722e636f6d2f52557469567a482e706e67\">\n\n> I will be integrating W&B for visualizations and logging artifacts!\n> \n> [GatedTabTransformer in FLAX project on W&B Dashboard](https://wandb.ai/usharengaraju/GatedTabTransformer_FLAX)\n> \n> - To get the API key, create an account in the [website](https://wandb.ai/site) .\n> - Use secrets to use API Keys more securely","metadata":{}},{"cell_type":"markdown","source":"# **<span style=\"color:#F7B2B0;\">W & B Artifacts</span>**\n\nAn artifact as a versioned folder of data.Entire datasets can be directly stored as artifacts .\n\nW&B Artifacts are used for dataset versioning, model versioning . They are also used for tracking dependencies and results across machine learning pipelines.Artifact references can be used to point to data in other systems like S3, GCP, or your own system.\n\nYou can learn more about W&B artifacts [here](https://docs.wandb.ai/guides/artifacts)\n\n![](https://drive.google.com/uc?id=1JYSaIMXuEVBheP15xxuaex-32yzxgglV)","metadata":{}},{"cell_type":"code","source":"import wandb\nfrom kaggle_secrets import UserSecretsClient\nuser_secrets = UserSecretsClient()\nwandb_key = user_secrets.get_secret(\"api_key\")\nwandb.login(key = wandb_key)","metadata":{"execution":{"iopub.status.busy":"2022-11-25T16:31:07.483378Z","iopub.execute_input":"2022-11-25T16:31:07.483954Z","iopub.status.idle":"2022-11-25T16:31:10.831783Z","shell.execute_reply.started":"2022-11-25T16:31:07.483834Z","shell.execute_reply":"2022-11-25T16:31:10.830710Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Save training data to W&B Artifacts\nrun = wandb.init(project='GatedTabTransformer_FLAX', name='processed_data') \nartifact = wandb.Artifact(name='processed_data',type='dataset')\nartifact.add_file(\"/kaggle/input/amex-tfrecords/minidata (1).csv\")\nwandb.log_artifact(artifact)\nwandb.finish()","metadata":{"execution":{"iopub.status.busy":"2022-11-25T16:31:14.663199Z","iopub.execute_input":"2022-11-25T16:31:14.663644Z","iopub.status.idle":"2022-11-25T16:32:08.582818Z","shell.execute_reply.started":"2022-11-25T16:31:14.663601Z","shell.execute_reply":"2022-11-25T16:32:08.581615Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"![](https://i.imgur.com/tRDVISy.png)","metadata":{}},{"cell_type":"markdown","source":" # <h1 style='background:#F7B2B0; border:0; color:black'><center>Model Architecture</center></h1> ","metadata":{}},{"cell_type":"code","source":"class Sequential(Module):\n    layers: Sequence[Type[Module]]\n\n    @nn.compact\n    def __call__(self, x):\n        for layer in self.layers:\n            x = layer(x)\n        return x\n\n\nclass Residual(Module):\n    layers: Sequence[Type[Module]]\n    @nn.compact\n    def __call__(self, x):\n        for layer in self.layers:\n            x = layer(x) + x\n        return x\n\n\nclass PreNorm(Module):\n    layers: Sequence[Type[Module]]\n    def setup(self):\n        self.norm = nn.LayerNorm()\n    @nn.compact\n    def __call__(self, x):\n        for layer in self.layers:\n            x = self.norm(x)\n            x = layer(x)\n        return x\n\nclass Identity(Module):\n    @nn.compact\n    def __call__(self, x):\n        return x\n\nclass Attention(nn.Module):\n    dim_out: int\n    dim_head: int\n    dtype: Dtype = jnp.float32\n    def setup(self):\n        self.scale = self.dim_head ** -0.5\n        self.to_qkv = nn.Dense(features=self.dim_head * 3, dtype=self.dtype)\n        self.to_out = nn.Dense(features=self.dim_out, dtype=self.dtype)\n    @nn.compact\n    def __call__(self, x) -> Array:\n        n = x.shape[0]\n        qkv = self.to_qkv(x)\n        q, k, v = jnp.split(qkv, 3, axis=-1)\n        sim = jnp.einsum(\"i d, j d -> i j\", q, k) * self.scale\n        mask = jnp.triu(jnp.ones((n, n), dtype=bool), 1)\n        sim = jnp.where(mask, ATTN_MASK_VALUE, sim)\n        attn = nn.softmax(sim, axis=-1)\n        out = jnp.einsum(\"i j, j d -> i d\", attn, v)\n        return self.to_out(out)\n\nclass SpatialGatingUnit(nn.Module):\n    dim_out: int\n    dtype: Dtype = jnp.float32\n    def setup(self):\n        self.norm = LayerNorm(dtype=self.dtype)\n        self.proj_out = nn.Dense(features=self.dim_out, dtype=self.dtype)\n    @nn.compact\n    def __call__(self, x, gate_res=None) -> Array:\n        x, gate = jnp.split(x, 2, axis=-1)\n        gate = self.norm(gate)\n        if gate_res is not None:\n            gate += gate_res\n        x = x * gate\n        return self.proj_out(x)\n  \nclass gMLPBlock(nn.Module):\n    dim: int\n    dim_ff: int\n    attn_dim: Any = None\n    dtype: Dtype = jnp.float32\n\n    def setup(self):\n        self.proj_in = nn.Dense(features=self.dim_ff, dtype=self.dtype)\n        self.attn = (\n            Attention(\n                dim_head=self.attn_dim, dim_out=self.dim_ff // 2, dtype=self.dtype\n            )\n            if self.attn_dim is not None\n            else None\n        )\n        self.sgu = SpatialGatingUnit(dim_out=self.dim_ff // 2, dtype=self.dtype)\n        self.proj_out = nn.Dense(features=self.dim, dtype=self.dtype)\n\n    @nn.compact\n    def __call__(self, x) -> Array:\n        gate_res = self.attn(x) if self.attn is not None else None\n\n        x = self.proj_in(x)\n        x = nn.gelu(x)\n        x = self.sgu(x, gate_res=gate_res)\n        x = self.proj_out(x)\n        return x\n\n\nclass gMLP(nn.Module):\n\n    dim: int\n    depth: int\n    num_tokens: Any = None\n    ff_mult: int = 4\n    attn_dim: Any = None\n    dtype: Dtype = jnp.float32\n\n    def setup(self):\n        dim_ff = self.dim * self.ff_mult\n        self.to_embed = (\n            nn.Embed(\n                num_embeddings=self.num_tokens, features=self.dim, dtype=self.dtype\n            )\n            if self.num_tokens is not None\n            else Identity()\n        )\n\n        self.layers = [\n            Residual(\n                [\n                    PreNorm(\n                        [\n                            gMLPBlock(\n                                dim=self.dim,\n                                dim_ff=dim_ff,\n                                attn_dim=self.attn_dim,\n                                dtype=self.dtype,\n                            )\n                        ]\n                    )\n                ]\n            )\n            for i in range(self.depth)\n        ]\n\n        self.to_logits = (\n            Sequential(\n                [nn.LayerNorm(), nn.Dense(features=self.num_tokens, dtype=self.dtype)]\n            )\n            if self.num_tokens is not None\n            else Identity()\n        )\n\n    @nn.compact\n    def __call__(self, x) -> Array:\n        x = self.to_embed(x)\n        out = Sequential(self.layers)(x)\n        return self.to_logits(out)","metadata":{"id":"HNGPEUy5OB8r","execution":{"iopub.status.busy":"2022-11-21T18:49:01.541098Z","iopub.execute_input":"2022-11-21T18:49:01.541565Z","iopub.status.idle":"2022-11-21T18:49:01.576266Z","shell.execute_reply.started":"2022-11-21T18:49:01.541522Z","shell.execute_reply":"2022-11-21T18:49:01.574917Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def scaled_dot_product(q, k, v, mask=None):\n    d_k = q.shape[-1]\n    attn_logits = jnp.matmul(q, jnp.swapaxes(k, -2, -1))\n    attn_logits = attn_logits / math.sqrt(d_k)\n    # print('dot1',attn_logits[0], attn_logits.shape, jnp.amax(attn_logits),d_k)\n    attention = jax.nn.softmax(attn_logits)\n    # print('dot2',attention[0])\n    values = jnp.matmul(attention, v)\n    return values, attention\n\nclass MultiheadAttention(nn.Module):\n    embed_dim : int  # Output dimension\n    num_heads : int  # Number of parallel heads (h)\n    \n    def setup(self):\n        # Stack all weight matrices 1...h and W^Q, W^K, W^V together for efficiency\n        # Note that in many implementations you see \"bias=False\" which is optional\n        self.qkv_proj = nn.Dense(3*self.embed_dim,\n                                 kernel_init=nn.initializers.xavier_uniform(),  # Weights with Xavier uniform init\n                                 bias_init=nn.initializers.zeros  # Bias init with zeros\n                                )\n        self.o_proj = nn.Dense(self.embed_dim,\n                               kernel_init=nn.initializers.xavier_uniform(),\n                               bias_init=nn.initializers.zeros)\n\n    def __call__(self, x, mask=None):\n        batch_size, seq_length, embed_dim = x.shape\n        qkv = self.qkv_proj(x)\n        # Separate Q, K, V from linear output\n        qkv = qkv.reshape(batch_size, seq_length, self.num_heads, -1)\n        qkv = qkv.transpose(0, 2, 1, 3) # [Batch, Head, SeqLen, Dims]\n        q, k, v = jnp.array_split(qkv, 3, axis=-1)\n        # print('att0',jnp.isnan(jax.device_get(x)).any()==True )\n        # Determine value outputs\n        values, attention = scaled_dot_product(q, k, v, mask=mask)\n        # print('att1',jnp.isnan(jax.device_get(values)).any()==True )\n        values = values.transpose(0, 2, 1, 3) # [Batch, SeqLen, Head, Dims]\n        values = values.reshape(batch_size, seq_length, embed_dim)\n        o = self.o_proj(values)\n        return o, attention\n\nclass EncoderBlock(nn.Module):\n    input_dim : int  # Input dimension is needed here since it is equal to the output dimension (residual connection)\n    num_heads : int\n    dim_feedforward : int\n    dropout_prob : float\n    \n    def setup(self):\n        # Attention layer\n        self.self_attn = MultiheadAttention(embed_dim=self.input_dim, \n                                            num_heads=self.num_heads)\n        # Two-layer MLP\n        self.linear = [\n            nn.Dense(self.dim_feedforward),\n            nn.Dropout(self.dropout_prob),\n            nn.relu,\n            nn.Dense(self.input_dim)\n        ]\n        # Layers to apply in between the main layers\n        self.norm1 = nn.LayerNorm()\n        self.norm2 = nn.LayerNorm()\n        self.dropout = nn.Dropout(self.dropout_prob)\n\n    def __call__(self, x, mask=None, train=True):\n        # Attention part\n        # print('block0',jnp.isnan(jax.device_get(x)).any()==True )\n        attn_out, _ = self.self_attn(x, mask=mask)\n        # print('block-1',attn_out[0])\n        x = x + self.dropout(attn_out, deterministic=not train)\n        x = self.norm1(x)\n        # MLP part\n        linear_out = x\n        for l in self.linear:\n            linear_out = l(linear_out) if not isinstance(l, nn.Dropout) else l(linear_out, deterministic=not train)\n        x = x + self.dropout(linear_out, deterministic=not train)\n        x = self.norm2(x)\n        return x\n\nclass TransformerEncoder(nn.Module):\n    num_layers : int\n    input_dim : int\n    num_heads : int\n    dim_feedforward : int\n    dropout_prob : float\n    \n    def setup(self):\n        self.layers = [EncoderBlock(self.input_dim, self.num_heads, self.dim_feedforward, self.dropout_prob) for _ in range(self.num_layers)]\n\n    def __call__(self, x, mask=None, train=True):\n        for l in self.layers:\n            # print('transenc',x[0])\n            x = l(x, mask=mask, train=train)\n        # print('transenc',x[0])\n        return x\n\nclass GatedTabTransformer(nn.Module):\n    model_dim : int                   # Hidden dimensionality to use inside the Transformer\n    num_classes : int                 # Number of classes to predict per sequence element\n    num_heads : int                   # Number of heads to use in the Multi-Head Attention blocks\n    num_layers : int                  # Number of encoder blocks to use\n    categories : Sequence\n    gmlp_blocks : int\n    dropout_prob : float = 0.2        # Dropout to apply inside the model\n    input_dropout_prob : float = 0.0  # Dropout to apply on the input features\n    embedding_dim : int =256\n\n    def setup(self):\n        # Input dim -> Model dim\n        self.input_layer = nn.Dense(self.model_dim)\n        self.contnorm = nn.LayerNorm()\n        # Transformer\n        self.transformer = TransformerEncoder(num_layers=self.num_layers,\n                                              input_dim=self.model_dim,\n                                              dim_feedforward=2*self.model_dim,\n                                              num_heads=self.num_heads,\n                                              dropout_prob=self.dropout_prob)\n        # Output classifier per sequence lement\n        self.embed_layers = [nn.Embed(num_embeddings = 10, features = self.model_dim ) for number_of_classes in self.categories]\n        self.embedder2 = nn.Dense(self.embedding_dim)\n        # print(type(self.embed_layers))\n        # print(self.categories)\n        # for number_of_classes in self.categories:\n        #     self.embed_layers.append(nn.Embed(num_embeddings = self.model_dim, features = number_of_classes))\n        # print(22222)\n        self.gmlp_layers = []\n        self.gmlp_layer = gMLP(dim = self.embedding_dim , depth =self.gmlp_blocks)\n        self.output_net = [\n            nn.Dense(self.model_dim),\n            nn.LayerNorm(),\n            nn.relu,\n            nn.Dropout(self.dropout_prob),\n            nn.Dense(self.num_classes),\n            nn.softmax\n        ]\n\n    def __call__(self, cat_x,cont_x, mask=None,train=True):\n        cont_x = self.contnorm(cont_x)\n        embedding_outputs = []\n        for categorical_input, embedding_layer in zip(cat_x.T, self.embed_layers):\n            embedding_outputs.append(jnp.expand_dims(embedding_layer(categorical_input),axis=1))\n        # print(cat_x.shape)\n        categorical_inputs = jnp.concatenate(embedding_outputs,axis=1)\n        # print('cat',jnp.isnan(jax.device_get(cat_x)).any()==True )\n        # print(categorical_inputs.shape)\n        # print(1111, categorical_inputs.shape)\n        # categorical_inputs = self.input_layer(categorical_inputs)\n        # print(2222, categorical_inputs.shape)\n        categorical_inputs = self.transformer(categorical_inputs, mask=mask, train=train)\n        # print(categorical_inputs)\n        x,y,z = categorical_inputs.shape\n        # print(4444 , cont_x.shape)\n        # print(categorical_inputs.shape)\n        categorical_inputs = jnp.reshape(categorical_inputs,newshape=(x,y*z))\n        gmlp_inp = jnp.concatenate([categorical_inputs,cont_x],axis=1)\n        \n        gmlp_inp = jnp.expand_dims(self.embedder2(gmlp_inp),axis=1)\n        gmlp_inp = self.gmlp_layer(gmlp_inp)\n        \n        x = jnp.mean(gmlp_inp,axis=1)\n        for l in self.output_net:\n            x = l(x) if not isinstance(l, nn.Dropout) else l(x, deterministic=not train)\n        return x\n","metadata":{"id":"cFNafRwwP4ZL","execution":{"iopub.status.busy":"2022-11-21T18:49:01.578216Z","iopub.execute_input":"2022-11-21T18:49:01.578953Z","iopub.status.idle":"2022-11-21T18:49:01.614198Z","shell.execute_reply.started":"2022-11-21T18:49:01.578915Z","shell.execute_reply":"2022-11-21T18:49:01.613229Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"","metadata":{}},{"cell_type":"code","source":"import math\nmain_rng = random.PRNGKey(42)\nmain_rng, x_rng = random.split(main_rng)\n\n\ntranspre = GatedTabTransformer(num_layers=5, \n                                model_dim=128,\n                                num_classes=2,\n                                num_heads=4,\n                                dropout_prob=0.15,\n                                input_dropout_prob=0.05,\n                                gmlp_blocks = 6,\n                                categories = cats)\n# Initialize parameters of transformer predictor with random key and inputs\nmain_rng, init_rng, dropout_init_rng = random.split(main_rng, 3)\nparams = transpre.init({'params': init_rng, 'dropout': dropout_init_rng}, cat_arr,cont_arr, train=True)['params']\n# Apply transformer predictor with parameters on the inputs\n# Since dropout is stochastic, we need to pass a rng to the forward\nmain_rng, dropout_apply_rng = random.split(main_rng)\n# Instead of passing params and rngs every time to a function call, we can bind them to the module\nbinded_mod = transpre.bind({'params': params}, rngs={'dropout': dropout_apply_rng})\nout = binded_mod(cat_arr,cont_arr, train=True)\nprint('Out', out.shape)","metadata":{"id":"GTKgleNgFO-a","outputId":"c577b150-5c05-44a2-fbe2-9ce4997d36f4","execution":{"iopub.status.busy":"2022-11-21T18:49:01.615628Z","iopub.execute_input":"2022-11-21T18:49:01.616809Z","iopub.status.idle":"2022-11-21T18:49:06.465805Z","shell.execute_reply.started":"2022-11-21T18:49:01.616762Z","shell.execute_reply":"2022-11-21T18:49:06.463628Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":" # <h1 style='background:#F7B2B0; border:0; color:black'><center>Model Training</center></h1> ","metadata":{}},{"cell_type":"code","source":"\ny_arr1 = jax.nn.one_hot(y_arr, num_classes=2)\nfrom flax.training import train_state  # Useful dataclass to keep train state\nfrom tqdm import tqdm\nimport numpy as np                     # Ordinary NumPy\nimport optax  \n\ndef cross_entropy_loss(*, logits, labels):\n  # labels_onehot = jax.nn.one_hot(labels, num_classes=10)\n  return optax.softmax_cross_entropy(logits=logits, labels=labels).mean()\n\ndef compute_metrics(*, logits, labels):\n  loss = cross_entropy_loss(logits=logits, labels=labels)\n  # print(logits.shape, labels.shape)\n  arr = jnp.argmax(logits, -1)\n  # print(arr.shape)\n  accuracy = jnp.mean(jnp.argmax(logits, -1) == jnp.argmax(labels, -1))\n  metrics = {\n      'loss': loss,\n      'accuracy': accuracy,\n  }\n  return metrics\ndef create_train_state(init_rng,dropout_init_rng,cat,cont, learning_rate, momentum):\n  \"\"\"Creates initial `TrainState`.\"\"\"\n  model = GatedTabTransformer(num_layers=5, \n                                model_dim=128,\n                                num_classes=2,\n                                num_heads=4,\n                                dropout_prob=0.15,\n                                input_dropout_prob=0.05,\n                                gmlp_blocks = 6,\n                                categories = cats)\n  params = model.init({'params': init_rng, 'dropout': dropout_init_rng}, cat,cont, train=True)['params']\n  tx = optax.sgd(learning_rate, momentum)\n  return train_state.TrainState.create(\n      apply_fn=model.apply, params=params, tx=tx)\n\n@jax.jit\ndef train_step(state, cat,cont,y,dropout_apply_rng):\n  \"\"\"Train for a single step.\"\"\"\n  def loss_fn(params):\n    model = GatedTabTransformer(num_layers=5, \n                                model_dim=128,\n                                num_classes=2,\n                                num_heads=4,\n                                dropout_prob=0.15,\n                                input_dropout_prob=0.05,\n                                gmlp_blocks = 6,\n                                categories = cats)\n    logits = model.apply({'params': params}, cat,cont,rngs={'dropout': dropout_apply_rng})\n    loss = cross_entropy_loss(logits=logits, labels=y)\n    return loss, logits\n  grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n  (_, logits), grads = grad_fn(state.params)\n  state = state.apply_gradients(grads=grads)\n  metrics = compute_metrics(logits=logits, labels=y)\n  return state, metrics\n\n# def train_epoch(state, cat,cont,y, epoch, rng):\n#   main_rng, dropout_apply_rng = random.split(rng)\n  \n#   state, metrics = train_step(state, cat,cont,y,dropout_apply_rng)\n#   print('train epoch: %d, loss: %.4f, accuracy: %.2f' % (\n#       epoch, metrics['loss'], metrics['accuracy'] * 100))\n\n#   return state\n\ndef train_epoch(state, cat,cont,y, epoch, rng):\n  main_rng, dropout_apply_rng = random.split(rng)\n  batch_metrics = []\n  for i in tqdm(range(0,30000,32)):\n    cat_x = cat[i:i+32]    \n    cont_x = cont[i:i+32]\n    yy = y[i:i+32]\n    state, metrics = train_step(state, cat_x,cont_x,yy,dropout_apply_rng)\n    batch_metrics.append(metrics)\n\n  batch_metrics_np = jax.device_get(batch_metrics)\n  epoch_metrics_np = {\n      k: np.mean([metrics[k] for metrics in batch_metrics_np])\n      for k in batch_metrics_np[0]}\n\n  print('train epoch: %d, loss: %.4f, accuracy: %.2f' % (\n      epoch, epoch_metrics_np['loss'], epoch_metrics_np['accuracy'] * 100))\n\n  return state,epoch_metrics_np\n\nlearning_rate = 0.001\nmomentum = 0.9\nmain_rng = random.PRNGKey(42)\nmain_rng, x_rng = random.split(main_rng)\nstate = create_train_state(main_rng,x_rng,cat_arr,cont_arr, learning_rate, momentum)\nall_mets=[]\nfor epoch in range(1, 10 + 1):\n  # Use a separate PRNG key to permute image data during shuffling\n  rng, input_rng = jax.random.split(main_rng)\n  # Run an optimization step over a training batch\n  state,mets = train_epoch(state, cat_arr, cont_arr,y_arr1, epoch, input_rng)\n  all_mets.append(mets)","metadata":{"id":"MDir0dm_OjH9","outputId":"04f61bcf-a6ca-4876-d45f-13ff7d18d3b6","execution":{"iopub.status.busy":"2022-11-21T18:49:06.469616Z","iopub.status.idle":"2022-11-21T18:49:06.470629Z","shell.execute_reply.started":"2022-11-21T18:49:06.470318Z","shell.execute_reply":"2022-11-21T18:49:06.470349Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":" # <h1 style='background:#F7B2B0; border:0; color:black'><center>Visualizing the Metrics</center></h1> ","metadata":{}},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n\n# summarize history for accuracy\nplt.plot([mets['accuracy'] for mets in all_mets])\nplt.title('model accuracy')\nplt.ylabel('accuracy')\nplt.xlabel('epoch')\nplt.legend(['train', 'test'], loc='upper left')\nplt.show()\n\n\n# summarize history for loss\nplt.plot([mets['loss'] for mets in all_mets])\nplt.title('model loss')\nplt.ylabel('loss')\nplt.xlabel('epoch')\nplt.legend(['train', 'test'], loc='upper left')\nplt.show()","metadata":{"id":"c4_oBirWpBFK","outputId":"20bc7404-5a80-46e9-dc4d-1fdc514db05e","execution":{"iopub.status.busy":"2022-11-21T18:49:06.471751Z","iopub.status.idle":"2022-11-21T18:49:06.472498Z","shell.execute_reply.started":"2022-11-21T18:49:06.472281Z","shell.execute_reply":"2022-11-21T18:49:06.472309Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"TabTransformers are highly robust against missing and noisy data and provide better interpretability . By replacing the final MLP block with gated Multilayer perceptron , GatedTabTranformers are able to achieve high accuracy in binary classification tasks.\n\n### References\n\nhttps://arxiv.org/pdf/2201.00199.pdf\n\nhttps://arxiv.org/pdf/2012.06678.pdf\n\nhttps://www.tensorflow.org/\n\nhttps://flax.readthedocs.io/\n\nPytorch Implementation : https://github.com/radi-cho/GatedTabTransformer","metadata":{}}]}