-- created Tab Transformer config

manujosephv · manujosephv · commit 1c8bdabb6c72 · 2021-06-20T08:28:43.000+05:30
diff --git a/pytorch_tabular/models/tab_transformer/__init__.py b/pytorch_tabular/models/tab_transformer/__init__.py
@@ -0,0 +1,4 @@
+from .autoint import AutoIntBackbone, AutoIntModel
+from .config import AutoIntConfig
+
+__all__ = ["AutoIntModel", "AutoIntBackbone", "AutoIntConfig"]
diff --git a/pytorch_tabular/models/tab_transformer/config.py b/pytorch_tabular/models/tab_transformer/config.py
@@ -0,0 +1,124 @@
+# Pytorch Tabular
+# Author: Manu Joseph <manujoseph@gmail.com>
+# For license information, see LICENSE.TXT
+"""AutomaticFeatureInteraction Config"""
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+from pytorch_tabular.config import ModelConfig, _validate_choices
+
+
+@dataclass
+class TabTransformerConfig(ModelConfig):
+    """Tab Transformer configuration
+    Args:
+        task (str): Specify whether the problem is regression of classification.Choices are: regression classification
+        learning_rate (float): The learning rate of the model
+        loss (Union[str, NoneType]): The loss function to be applied.
+            By Default it is MSELoss for regression and CrossEntropyLoss for classification.
+            Unless you are sure what you are doing, leave it at MSELoss or L1Loss for regression and CrossEntropyLoss for classification
+        metrics (Union[List[str], NoneType]): the list of metrics you need to track during training.
+            The metrics should be one of the metrics implemented in PyTorch Lightning.
+            By default, it is Accuracy if classification and MeanSquaredLogError for regression
+        metrics_params (Union[List, NoneType]): The parameters to be passed to the Metrics initialized
+        target_range (Union[List, NoneType]): The range in which we should limit the output variable. Currently ignored for multi-target regression
+            Typically used for Regression problems. If left empty, will not apply any restrictions
+
+        attn_embed_dim (int): The number of hidden units in the Multi-Headed Attention layers. Defaults to 32
+        num_heads (int): The number of heads in the Multi-Headed Attention layer. Defaults to 2
+        num_attn_blocks (int): The number of layers of stacked Multi-Headed Attention layers. Defaults to 2
+        attn_dropouts (float): Dropout between layers of Multi-Headed Attention Layers. Defaults to 0.0
+        has_residuals (bool): Flag to have a residual connect from enbedded output to attention layer output.
+            Defaults to True
+        embedding_dim (int): The dimensions of the embedding for continuous and categorical columns. Defaults to 16
+        embedding_dropout (float): probability of an embedding element to be zeroed. Defaults to 0.0
+        deep_layers (bool): Flag to enable a deep MLP layer before the Multi-Headed Attention layer. Defaults to False
+        layers (str): Hyphen-separated number of layers and units in the deep MLP. Defaults to 128-64-32
+        activation (str): The activation type in the deep MLP. The default activaion in PyTorch like
+            ReLU, TanH, LeakyReLU, etc. https://pytorch.org/docs/stable/nn.html#non-linear-activations-weighted-sum-nonlinearity.
+            Defaults to ReLU
+        dropout (float): probability of an classification element to be zeroed in the deep MLP. Defaults to 0.0
+        use_batch_norm (bool): Flag to include a BatchNorm layer after each Linear Layer+DropOut. Defaults to False
+        batch_norm_continuous_input (bool): If True, we will normalize the contiinuous layer by passing it through a BatchNorm layer. Defaults to False
+        attention_pooling (bool): If True, will combine the attention outputs of each block for final prediction. Defaults to False
+        initialization (str): Initialization scheme for the linear layers. Defaults to `kaiming`.
+            Choices are: [`kaiming`,`xavier`,`random`].
+
+    Raises:
+        NotImplementedError: Raises an error if task is not in ['regression','classification']
+    """
+
+    transformer_embed_dim: int = field(
+        default=32,
+        metadata={
+            "help": "The number of hidden units in the Multi-Headed Attention layers. Defaults to 32"
+        },
+    )
+    num_heads: int = field(
+        default=8,
+        metadata={
+            "help": "The number of heads in the Multi-Headed Attention layer. Defaults to 8"
+        },
+    )
+    num_attn_blocks: int = field(
+        default=6,
+        metadata={
+            "help": "The number of layers of stacked Multi-Headed Attention layers. Defaults to 6"
+        },
+    )
+    attn_dropouts: float = field(
+        default=0.1,
+        metadata={
+            "help": "Dropout between layers of Multi-Headed Attention Layers. Defaults to 0.1"
+        },
+    )
+    ff_dropouts: float = field(
+        default=0.1,
+        metadata={
+            "help": "Dropout after FF layers. Defaults to 0.1"
+        },
+    )
+    ff_hidden_multipliers: tuple = field(
+        default=(4,2),
+        metadata={
+            "help": "Multiples by which the layers scale from Transformer output to logits. Defaults to (4,2)"
+        },
+    )
+    mlp_activation: str = field(
+        default="ReLU",
+        metadata={
+            "help": "The activation type in the final FF layer. The default activaion in PyTorch like ReLU, TanH, LeakyReLU, etc. https://pytorch.org/docs/stable/nn.html#non-linear-activations-weighted-sum-nonlinearity. Defaults to ReLU"
+        },
+    )
+    transformer_activation: str = field(
+        default="GEGLU",
+        metadata={
+            "help": "The activation type in the transformer feed forward layers. In addition to the default activation in PyTorch like ReLU, TanH, LeakyReLU, etc. https://pytorch.org/docs/stable/nn.html#non-linear-activations-weighted-sum-nonlinearity, GatedGLUs are also implemented(https://arxiv.org/pdf/2002.05202.pdf). Defaults to GEGLU"
+        },
+    )
+    initialization: str = field(
+        default="kaiming",
+        metadata={
+            "help": "Initialization scheme for the linear layers. Defaults to `kaiming`",
+            "choices": ["kaiming", "xavier", "random"],
+        },
+    )
+    _module_src: str = field(default="tab_transformer")
+    _model_name: str = field(default="TabTransformerModel")
+    _config_name: str = field(default="TabTransformerConfig")
+
+
+# cls = AutoIntConfig
+# desc = "Configuration for Data."
+# doc_str = f"{desc}\nArgs:"
+# for key in cls.__dataclass_fields__.keys():
+#     atr = cls.__dataclass_fields__[key]
+#     if atr.init:
+#         type = str(atr.type).replace("<class '","").replace("'>","").replace("typing.","")
+#         help_str = atr.metadata.get("help","")
+#         if "choices" in atr.metadata.keys():
+#             help_str += f'. Choices are: [{",".join(["`"+str(ch)+"`" for ch in atr.metadata["choices"]])}].'
+#         # help_str += f'. Defaults to {atr.default}'
+#         doc_str+=f'\n\t\t{key} ({type}): {help_str}'
+
+# print(doc_str)
diff --git a/pytorch_tabular/models/tab_transformer/tab_transformer.py b/pytorch_tabular/models/tab_transformer/tab_transformer.py
@@ -0,0 +1,163 @@
+# Pytorch Tabular
+# Author: Manu Joseph <manujoseph@gmail.com>
+# For license information, see LICENSE.TXT
+# Inspired by https://github.com/lucidrains/tab-transformer-pytorch/blob/main/tab_transformer_pytorch/tab_transformer_pytorch.py
+"""TabTransformer Model"""
+import logging
+from typing import Dict
+
+import pytorch_lightning as pl
+import torch
+import torch.nn as nn
+from omegaconf import DictConfig
+
+from pytorch_tabular.utils import _initialize_layers, _linear_dropout_bn
+
+from ..base_model import BaseModel
+
+logger = logging.getLogger(__name__)
+
+#TODO dont use embedding_dims
+class TabTransformerBackbone(pl.LightningModule):
+    def __init__(self, config: DictConfig):
+        super().__init__()
+        self.save_hyperparameters(config)
+        self._build_network()
+
+    def _build_network(self):
+        if len(self.hparams.categorical_cols)>0:
+            # Category Embedding layers
+            self.cat_embedding_layers = nn.ModuleList(
+                [
+                    nn.Embedding(cardinality, self.hparams.embedding_dim)
+                    for cardinality in self.hparams.categorical_cardinality
+                ]
+            )
+        if self.hparams.batch_norm_continuous_input:
+            self.normalizing_batch_norm = nn.BatchNorm1d(self.hparams.continuous_dim)
+        # Continuous Embedding Layer
+        self.cont_embedding_layer = nn.Embedding(
+            self.hparams.continuous_dim, self.hparams.embedding_dim
+        )
+        if self.hparams.embedding_dropout != 0 and len(self.hparams.categorical_cols)>0:
+            self.embed_dropout = nn.Dropout(self.hparams.embedding_dropout)
+        # Deep Layers
+        _curr_units = self.hparams.embedding_dim
+        if self.hparams.deep_layers:
+            activation = getattr(nn, self.hparams.activation)
+            # Linear Layers
+            layers = []
+            for units in self.hparams.layers.split("-"):
+                layers.extend(
+                    _linear_dropout_bn(
+                        self.hparams,
+                        _curr_units,
+                        int(units),
+                        activation,
+                        self.hparams.dropout,
+                    )
+                )
+                _curr_units = int(units)
+            self.linear_layers = nn.Sequential(*layers)
+        # Projection to Multi-Headed Attention Dims
+        self.attn_proj = nn.Linear(_curr_units, self.hparams.attn_embed_dim)
+        _initialize_layers(self.hparams, self.attn_proj)
+        # Multi-Headed Attention Layers
+        self.self_attns = nn.ModuleList(
+            [
+                nn.MultiheadAttention(
+                    self.hparams.attn_embed_dim,
+                    self.hparams.num_heads,
+                    dropout=self.hparams.attn_dropouts,
+                )
+                for _ in range(self.hparams.num_attn_blocks)
+            ]
+        )
+        if self.hparams.has_residuals:
+            self.V_res_embedding = torch.nn.Linear(
+                _curr_units,
+                self.hparams.attn_embed_dim * self.hparams.num_attn_blocks
+                if self.hparams.attention_pooling
+                else self.hparams.attn_embed_dim,
+            )
+        self.output_dim = (
+            self.hparams.continuous_dim + self.hparams.categorical_dim
+        ) * self.hparams.attn_embed_dim
+        if self.hparams.attention_pooling:
+            self.output_dim = self.output_dim * self.hparams.num_attn_blocks
+
+    def forward(self, x: Dict):
+        # (B, N)
+        continuous_data, categorical_data = x["continuous"], x["categorical"]
+        x = None
+        if len(self.hparams.categorical_cols) > 0:
+            x_cat = [
+                embedding_layer(categorical_data[:, i]).unsqueeze(1)
+                for i, embedding_layer in enumerate(self.cat_embedding_layers)
+            ]
+            # (B, N, E)
+            x = torch.cat(x_cat, 1)
+        if self.hparams.continuous_dim > 0:
+            cont_idx = (
+                torch.arange(self.hparams.continuous_dim)
+                .expand(continuous_data.size(0), -1)
+                .to(self.device)
+            )
+            if self.hparams.batch_norm_continuous_input:
+                continuous_data = self.normalizing_batch_norm(continuous_data)
+            x_cont = torch.mul(
+                continuous_data.unsqueeze(2),
+                self.cont_embedding_layer(cont_idx),
+            )
+            # (B, N, E)
+            x = x_cont if x is None else torch.cat([x, x_cont], 1)
+        if self.hparams.embedding_dropout != 0 and len(self.hparams.categorical_cols) > 0:
+            x = self.embed_dropout(x)
+        if self.hparams.deep_layers:
+            x = self.linear_layers(x)
+        # (N, B, E*) --> E* is the Attn Dimention
+        cross_term = self.attn_proj(x).transpose(0, 1)
+        if self.hparams.attention_pooling:
+            attention_ops = []
+        for self_attn in self.self_attns:
+            cross_term, _ = self_attn(cross_term, cross_term, cross_term)
+            if self.hparams.attention_pooling:
+                attention_ops.append(cross_term)
+        if self.hparams.attention_pooling:
+            cross_term = torch.cat(attention_ops, dim=-1)
+        # (B, N, E*)
+        cross_term = cross_term.transpose(0, 1)
+        if self.hparams.has_residuals:
+            # (B, N, E*) --> Projecting Embedded input to Attention sub-space
+            V_res = self.V_res_embedding(x)
+            cross_term = cross_term + V_res
+        # (B, NxE*)
+        cross_term = nn.ReLU()(cross_term).reshape(-1, self.output_dim)
+        return cross_term
+
+
+class TabTransformerModel(BaseModel):
+    def __init__(self, config: DictConfig, **kwargs):
+        super().__init__(config, **kwargs)
+
+    def _build_network(self):
+        # Backbone
+        self.backbone = TabTransformerBackbone(self.hparams)
+        self.dropout = nn.Dropout(self.hparams.dropout)
+        # Adding the last layer
+        self.output_layer = nn.Linear(
+            self.backbone.output_dim, self.hparams.output_dim
+        )  # output_dim auto-calculated from other config
+        _initialize_layers(self.hparams, self.output_layer)
+
+    def forward(self, x: Dict):
+        x = self.backbone(x)
+        x = self.dropout(x)
+        y_hat = self.output_layer(x)
+        if (self.hparams.task == "regression") and (
+            self.hparams.target_range is not None
+        ):
+            for i in range(self.hparams.output_dim):
+                y_min, y_max = self.hparams.target_range[i]
+                y_hat[:, i] = y_min + nn.Sigmoid()(y_hat[:, i]) * (y_max - y_min)
+        return {"logits": y_hat, "backbone_features": x}
