Add Built-in Support for Model Stacking

docs/apidocs_model.md

@@ -30,6 +30,9 @@
 ::: pytorch_tabular.models.TabTransformerConfig
     options:
             heading_level: 3
+::: pytorch_tabular.models.StackingModelConfig
+    options:
+            heading_level: 3
 ::: pytorch_tabular.config.ModelConfig
     options:
             heading_level: 3
@@ -66,7 +69,9 @@
 ::: pytorch_tabular.models.TabTransformerModel
     options:
             heading_level: 3
-
+::: pytorch_tabular.models.StackingModel
+    options:
+            heading_level: 3
 ## Base Model Class
 ::: pytorch_tabular.models.BaseModel
     options:

docs/models.md

@@ -253,6 +253,30 @@ All the parameters have beet set to recommended values from the paper. Let's loo
 **For a complete list of parameters refer to the API Docs**
 [pytorch_tabular.models.DANetConfig][]
 
+## Model Stacking
+
+Model stacking is an ensemble learning technique that combines multiple base models to create a more powerful predictive model. Each base model processes the input features independently, and their outputs are concatenated before making the final prediction. This allows the model to leverage different learning patterns captured by each backbone architecture. You can use it by choosing `StackingModelConfig`. 
+
+The following diagram shows the concept of model stacking in PyTorch Tabular.
+![Model Stacking](imgs/model_stacking_concept.png)
+
+The following model architectures are supported for stacking:
+- Category Embedding Model
+- TabNet Model
+- FTTransformer Model
+- Gated Additive Tree Ensemble Model
+- DANet Model
+- AutoInt Model
+- GANDALF Model
+- Node Model
+
+All the parameters have been set to provide flexibility while maintaining ease of use. Let's look at them:
+
+- `model_configs`: List[ModelConfig]: List of configurations for each base model. Each config should be a valid PyTorch Tabular model config (e.g., NodeConfig, GANDALFConfig)
+
+**For a complete list of parameters refer to the API Docs**
+[pytorch_tabular.models.StackingModelConfig][]
+
 ## Implementing New Architectures
 
 PyTorch Tabular is very easy to extend and infinitely customizable. All the models that have been implemented in PyTorch Tabular inherits an Abstract Class `BaseModel` which is in fact a PyTorchLightning Model.

mkdocs.yml

@@ -24,6 +24,7 @@ nav:
                 - SHAP, Deep LIFT and so on through Captum Integration: "tutorials/14-Explainability.ipynb"
           - Custom PyTorch Models:
                 - Implementing New Supervised Architectures: "tutorials/04-Implementing New Architectures.ipynb"
+          - Model Stacking: "tutorials/16-Model Stacking.ipynb"
           - Other Features:
                 - Using Neural Categorical Embeddings in Scikit-Learn Workflows: "tutorials/03-Neural Embedding in Scikit-Learn Workflows.ipynb"
                 - Self-Supervised Learning using Denoising Autoencoders: "tutorials/08-Self-Supervised Learning-DAE.ipynb"

src/pytorch_tabular/models/__init__.py

@@ -19,6 +19,7 @@
 from .gate import GatedAdditiveTreeEnsembleConfig, GatedAdditiveTreeEnsembleModel
 from .mixture_density import MDNConfig, MDNModel
 from .node import NodeConfig, NODEModel
+from .stacking import StackingModel, StackingModelConfig
 from .tab_transformer import TabTransformerConfig, TabTransformerModel
 from .tabnet import TabNetModel, TabNetModelConfig
 
@@ -45,6 +46,8 @@
     "GANDALFBackbone",
     "DANetConfig",
     "DANetModel",
+    "StackingModel",
+    "StackingModelConfig",
     "category_embedding",
     "node",
     "mixture_density",
@@ -55,4 +58,5 @@
     "gate",
     "gandalf",
     "danet",
+    "stacking",
 ]

src/pytorch_tabular/models/stacking/__init__.py

@@ -0,0 +1,4 @@
+from .config import StackingModelConfig
+from .stacking_model import StackingBackbone, StackingModel
+
+__all__ = ["StackingModel", "StackingModelConfig", "StackingBackbone"]

src/pytorch_tabular/models/stacking/config.py

@@ -0,0 +1,26 @@
+from dataclasses import dataclass, field
+
+from pytorch_tabular.config import ModelConfig
+
+
+@dataclass
+class StackingModelConfig(ModelConfig):
+    """StackingModelConfig is a configuration class for the StackingModel. It is used to stack multiple models
+    together. Now, CategoryEmbeddingModel, TabNetModel, FTTransformerModel, GatedAdditiveTreeEnsembleModel, DANetModel,
+    AutoIntModel, GANDALFModel, NodeModel are supported.
+
+    Args:
+        model_configs (list[ModelConfig]): List of model configs to stack.
+
+    """
+
+    model_configs: list = field(default_factory=list, metadata={"help": "List of model configs to stack"})
+    _module_src: str = field(default="models.stacking")
+    _model_name: str = field(default="StackingModel")
+    _backbone_name: str = field(default="StackingBackbone")
+    _config_name: str = field(default="StackingConfig")
+
+
+# if __name__ == "__main__":
+#    from pytorch_tabular.utils import generate_doc_dataclass
+#    print(generate_doc_dataclass(StackingModelConfig))

src/pytorch_tabular/models/stacking/stacking_model.py

@@ -0,0 +1,140 @@
+import inspect
+
+import torch
+import torch.nn as nn
+from omegaconf import DictConfig
+
+import pytorch_tabular.models as models
+from pytorch_tabular.models import BaseModel
+from pytorch_tabular.models.common.heads import blocks
+from pytorch_tabular.models.gate import GatedAdditiveTreesBackbone
+from pytorch_tabular.models.node import NODEBackbone
+
+
+def instantiate_backbone(hparams, backbone_name):
+    backbone_class = getattr(getattr(models, hparams._module_src.split(".")[-1]), backbone_name)
+    class_args = list(inspect.signature(backbone_class).parameters.keys())
+    if "config" in class_args:
+        return backbone_class(config=hparams)
+    else:
+        return backbone_class(
+            **{
+                arg: getattr(hparams, arg) if arg != "block_activation" else getattr(nn, getattr(hparams, arg))()
+                for arg in class_args
+            }
+        )
+
+
+class StackingEmbeddingLayer(nn.Module):
+    def __init__(self, embedding_layers: nn.ModuleList):
+        super().__init__()
+        self.embedding_layers = embedding_layers
+
+    def forward(self, x):
+        outputs = []
+        for embedding_layer in self.embedding_layers:
+            em_output = embedding_layer(x)
+            outputs.append(em_output)
+        return outputs
+
+
+class StackingBackbone(nn.Module):
+    def __init__(self, config: DictConfig):
+        super().__init__()
+        self.hparams = config
+        self._build_network()
+
+    def _build_network(self):
+        self._backbones = nn.ModuleList()
+        self._heads = nn.ModuleList()
+        self._backbone_output_dims = []
+        assert len(self.hparams.model_configs) > 0, "Stacking requires more than 0 model"
+        for model_i in range(len(self.hparams.model_configs)):
+            # move necessary params to each model config
+            self.hparams.model_configs[model_i].embedded_cat_dim = self.hparams.embedded_cat_dim
+            self.hparams.model_configs[model_i].continuous_dim = self.hparams.continuous_dim
+            self.hparams.model_configs[model_i].n_continuous_features = self.hparams.continuous_dim
+
+            self.hparams.model_configs[model_i].embedding_dims = self.hparams.embedding_dims
+            self.hparams.model_configs[model_i].categorical_cardinality = self.hparams.categorical_cardinality
+            self.hparams.model_configs[model_i].categorical_dim = self.hparams.categorical_dim
+            self.hparams.model_configs[model_i].cat_embedding_dims = self.hparams.embedding_dims
+
+            # if output_dim is not set, set it to 128
+            if getattr(self.hparams.model_configs[model_i], "output_dim", None) is None:
+                self.hparams.model_configs[model_i].output_dim = 128
+
+            # if inferred_config is not set, set it to None.
+            if getattr(self.hparams, "inferred_config", None) is not None:
+                self.hparams.model_configs[model_i].inferred_config = self.hparams.inferred_config
+
+            # instantiate backbone
+            _backbone = instantiate_backbone(
+                self.hparams.model_configs[model_i], self.hparams.model_configs[model_i]._backbone_name
+            )
+            # set continuous_dim
+            _backbone.continuous_dim = self.hparams.continuous_dim
+            # if output_dim is not set, set it to the output_dim in model_config
+            if getattr(_backbone, "output_dim", None) is None:
+                setattr(
+                    _backbone,
+                    "output_dim",
+                    self.hparams.model_configs[model_i].output_dim,
+                )
+            self._backbones.append(_backbone)
+            self._backbone_output_dims.append(_backbone.output_dim)
+
+        self.output_dim = sum(self._backbone_output_dims)
+
+    def _build_embedding_layer(self):
+        assert getattr(self, "_backbones", None) is not None, "Backbones are not built"
+        embedding_layers = nn.ModuleList()
+        for backbone in self._backbones:
+            if getattr(backbone, "_build_embedding_layer", None) is None:
+                embedding_layers.append(nn.Identity())
+            else:
+                embedding_layers.append(backbone._build_embedding_layer())
+        return StackingEmbeddingLayer(embedding_layers)
+
+    def forward(self, x_list):
+        outputs = []
+        for i, backbone in enumerate(self._backbones):
+            bb_output = backbone(x_list[i])
+            if len(bb_output.shape) == 3 and isinstance(backbone, GatedAdditiveTreesBackbone):
+                bb_output = bb_output.mean(dim=-1)
+            elif len(bb_output.shape) == 3 and isinstance(backbone, NODEBackbone):
+                bb_output = bb_output.mean(dim=1)
+            outputs.append(bb_output)
+        x = torch.cat(outputs, dim=1)
+        return x
+
+
+class StackingModel(BaseModel):
+    def __init__(self, config: DictConfig, **kwargs):
+        super().__init__(config, **kwargs)
+
+    def _build_network(self):
+        self._backbone = StackingBackbone(self.hparams)
+        self._embedding_layer = self._backbone._build_embedding_layer()
+        self.output_dim = self._backbone.output_dim
+        self._head = self._get_head_from_config()
+
+    def _get_head_from_config(self):
+        _head_callable = getattr(blocks, self.hparams.head)
+        return _head_callable(
+            in_units=self.output_dim,
+            output_dim=self.hparams.output_dim,
+            config=_head_callable._config_template(**self.hparams.head_config),
+        )
+
+    @property
+    def backbone(self):
+        return self._backbone
+
+    @property
+    def embedding_layer(self):
+        return self._embedding_layer
+
+    @property
+    def head(self):
+        return self._head

src/pytorch_tabular/models/tabnet/config.py

@@ -129,6 +129,7 @@ class TabNetModelConfig(ModelConfig):
     _module_src: str = field(default="models.tabnet")
     _model_name: str = field(default="TabNetModel")
     _config_name: str = field(default="TabNetModelConfig")
+    _backbone_name: str = field(default="TabNetBackbone")
 
 
 # if __name__ == "__main__":