-- added test cases for mdn and auto int

-- added autoint
-- refactored a few things

Author: manujosephv

examples/to_test_regression.py

@@ -12,6 +12,7 @@
 from pytorch_tabular.models.category_embedding.config import (
     CategoryEmbeddingModelConfig,
 )
+from pytorch_tabular.models import AutoIntModel, AutoIntConfig
 
 from pytorch_tabular.models.mixture_density import (
     CategoryEmbeddingMDNConfig, MixtureDensityHeadConfig, NODEMDNConfig
@@ -33,6 +34,8 @@
 dataset = fetch_california_housing(data_home="data", as_frame=True)
 dataset.frame["HouseAgeBin"] = pd.qcut(dataset.frame["HouseAge"], q=4)
 dataset.frame.HouseAgeBin = "age_" + dataset.frame.HouseAgeBin.cat.codes.astype(str)
+dataset.frame["AveRoomsBin"] = pd.qcut(dataset.frame["AveRooms"], q=3)
+dataset.frame.AveRoomsBin = "av_rm_" + dataset.frame.AveRoomsBin.cat.codes.astype(str)
 
 test_idx = dataset.frame.sample(int(0.2 * len(dataset.frame)), random_state=42).index
 test = dataset.frame[dataset.frame.index.isin(test_idx)]
@@ -49,7 +52,7 @@
         "Longitude",
     ],
     # continuous_cols=[],
-    categorical_cols=["HouseAgeBin"],
+    categorical_cols=["HouseAgeBin","AveRoomsBin"],
     continuous_feature_transform=None,  # "yeo-johnson",
     normalize_continuous_features=True,
 )
@@ -61,8 +64,9 @@
 #     mdn_config = mdn_config
 # )
 # # model_config.validate()
-model_config = NodeConfig(task="regression", depth=2, embed_categorical=False)
-trainer_config = TrainerConfig(checkpoints=None, max_epochs=5, gpus=1, profiler=None)
+# model_config = CategoryEmbeddingModelConfig(task="regression")
+model_config = AutoIntConfig(task="regression", deep_layers=True, embedding_dropout=0.2, batch_norm_continuous_input=True)
+trainer_config = TrainerConfig(checkpoints=None, max_epochs=25, gpus=1, profiler=None, fast_dev_run=False, auto_lr_find=True)
 # experiment_config = ExperimentConfig(
 #     project_name="DeepGMM_test",
 #     run_name="wand_debug",

pytorch_tabular/models/__init__.py

@@ -8,9 +8,12 @@
     MixtureDensityHeadConfig,
     NODEMDNConfig,
     NODEMDN,
+    AutoIntMDN,
+    AutoIntMDNConfig
 )
+from .autoint import AutoIntConfig, AutoIntModel
 from .base_model import BaseModel
-from . import category_embedding, node, mixture_density, tabnet
+from . import category_embedding, node, mixture_density, tabnet, autoint
 
 __all__ = [
     "CategoryEmbeddingModel",
@@ -26,8 +29,13 @@
     "MixtureDensityHeadConfig",
     "NODEMDNConfig",
     "NODEMDN",
+    "AutoIntMDN",
+    "AutoIntMDNConfig",
+    "AutoIntConfig",
+    "AutoIntModel",
     "category_embedding",
     "node",
     "mixture_density",
     "tabnet",
+    "autoint",
 ]

pytorch_tabular/models/autoint/__init__.py

@@ -1,4 +1,4 @@
-from .category_embedding_model import CategoryEmbeddingModel, FeedForwardBackbone
-from .config import CategoryEmbeddingModelConfig
+from .autoint import AutoIntBackbone, AutoIntModel
+from .config import AutoIntConfig
 
-__all__ = ["CategoryEmbeddingModel", "CategoryEmbeddingModelConfig", "FeedForwardBackbone"]
+__all__ = ["AutoIntModel", "AutoIntBackbone", "AutoIntConfig"]

pytorch_tabular/models/autoint/autoint.py

@@ -9,7 +9,8 @@
 import torch
 import torch.nn as nn
 from omegaconf import DictConfig
-from pytorch_tabular.utils import _initialize_layers
+
+from pytorch_tabular.utils import _initialize_layers, _linear_dropout_bn
 
 from ..base_model import BaseModel
 
@@ -21,40 +22,32 @@ def __init__(self, config: DictConfig, **kwargs):
         self.embedding_cat_dim = sum([y for x, y in config.embedding_dims])
         super().__init__(config, **kwargs)
 
-    def _linear_dropout_bn(self, in_units, out_units, activation, dropout):
-        layers = []
-        if self.hparams.use_batch_norm:
-            layers.append(nn.BatchNorm1d(num_features=in_units))
-        linear = nn.Linear(in_units, out_units)
-        _initialize_layers(self.hparams, linear)
-        layers.extend([linear, activation()])
-        if dropout != 0:
-            layers.append(nn.Dropout(dropout))
-        return layers
-
     def _build_network(self):
-        # Embedding layers
+        # Category Embedding layers
         self.cat_embedding_layers = nn.ModuleList(
-            [nn.Embedding(x, y) for x, y in self.hparams.cat_embedding_dims]
-        )
-        self.cont_embedding_layers = nn.ModuleList(
             [
-                nn.Embedding(1, self.hparams.cont_embedding_dim)
-                for i in range(self.hparams.continuous_dim)
+                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 self.embedding_cat_dim != 0:
             self.embed_dropout = nn.Dropout(self.hparams.embedding_dropout)
-        # if self.hparams.use_batch_norm:
-        #     self.normalizing_batch_norm = nn.BatchNorm1d(self.hparams.continuous_dim+self.hparams.embedding_cat_dim)
+        # Deep Layers
+        _curr_units = self.hparams.embedding_dim
         if self.hparams.deep_layers:
             activation = getattr(nn, self.hparams.activation)
             # Linear Layers
             layers = []
-            _curr_units = self.hparams.continuous_dim + self.embedding_cat_dim
             for units in self.hparams.layers.split("-"):
                 layers.extend(
-                    self._linear_dropout_bn(
+                    _linear_dropout_bn(
+                        self.hparams,
                         _curr_units,
                         int(units),
                         activation,
@@ -63,9 +56,10 @@ def _build_network(self):
                 )
                 _curr_units = int(units)
             self.linear_layers = nn.Sequential(*layers)
-        else:
-            _curr_units = self.hparams.continuous_dim + self.embedding_cat_dim
-
+        # 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(
@@ -76,15 +70,56 @@ def _build_network(self):
                 for _ in range(self.hparams.num_attn_blocks)
             ]
         )
-        self.atten_output_dim = (
-            len(self.hparams.continuous_cols + self.hparams.categorical_cols)
-            * self.hparams.atten_embed_dim
-        )
-
+        if self.hparams.has_residuals:
+            self.V_res_embedding = torch.nn.Linear(
+                _curr_units, self.hparams.attn_embed_dim
+            )
+        self.output_dim = (
+            self.hparams.continuous_dim + self.hparams.categorical_dim
+        ) * self.hparams.attn_embed_dim
 
-    def forward(self, x):
-        x = self.linear_layers(x)
-        return x
+    def forward(self, x: Dict):
+        # (B, N)
+        continuous_data, categorical_data = x["continuous"], x["categorical"]
+        x = None
+        if self.embedding_cat_dim != 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 self.embedding_cat_dim != 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)
+        for self_attn in self.self_attns:
+            cross_term, _ = self_attn(cross_term, cross_term, cross_term)
+        # (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 AutoIntModel(BaseModel):
@@ -94,46 +129,18 @@ def __init__(self, config: DictConfig, **kwargs):
         super().__init__(config, **kwargs)
 
     def _build_network(self):
-        # Embedding layers
-        self.embedding_layers = nn.ModuleList(
-            [nn.Embedding(x, y) for x, y in self.hparams.embedding_dims]
-        )
-        # Continuous Layers
-        if self.hparams.batch_norm_continuous_input:
-            self.normalizing_batch_norm = nn.BatchNorm1d(self.hparams.continuous_dim)
         # Backbone
         self.backbone = AutoIntBackbone(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 unpack_input(self, x: Dict):
-        continuous_data, categorical_data = x["continuous"], x["categorical"]
-        if self.embedding_cat_dim != 0:
-            x = []
-            # for i, embedding_layer in enumerate(self.embedding_layers):
-            #     x.append(embedding_layer(categorical_data[:, i]))
-            x = [
-                embedding_layer(categorical_data[:, i])
-                for i, embedding_layer in enumerate(self.embedding_layers)
-            ]
-            x = torch.cat(x, 1)
-
-        if self.hparams.continuous_dim != 0:
-            if self.hparams.batch_norm_continuous_input:
-                continuous_data = self.normalizing_batch_norm(continuous_data)
-
-            if self.embedding_cat_dim != 0:
-                x = torch.cat([x, continuous_data], 1)
-            else:
-                x = continuous_data
-        return x
-
     def forward(self, x: Dict):
-        x = self.unpack_input(x)
         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