Slight Improvements to GATE model (#213)

* enabled two more parameters to GATE model

* added a env setup shell script

* fixing tree attention

* Reduced complexity of initial parameters in GATE

Author: manujosephv

examples/__only_for_dev__/adhoc_scaffold.py

@@ -50,7 +50,7 @@ def print_metrics(y_true, y_pred, tag):
 
 from pytorch_tabular import TabularModel  # noqa: E402
 from pytorch_tabular.config import DataConfig, OptimizerConfig, TrainerConfig  # noqa: E402
-from pytorch_tabular.models import CategoryEmbeddingModelConfig  # noqa: E402
+from pytorch_tabular.models import GatedAdditiveTreeEnsembleConfig  # noqa: E402
 
 data_config = DataConfig(
     # target should always be a list. Multi-targets are only supported for regression.
@@ -68,10 +68,10 @@ def print_metrics(y_true, y_pred, tag):
     fast_dev_run=True,
 )
 optimizer_config = OptimizerConfig()
-model_config = CategoryEmbeddingModelConfig(
+model_config = GatedAdditiveTreeEnsembleConfig(
     task="classification",
-    # gflu_stages=3,
-    # tree_depth=2,
+    gflu_stages=3,
+    tree_depth=2,
     # layers="4096-4096-512",  # Number of nodes in each layer
     # activation="LeakyReLU",  # Activation between each layers
     learning_rate=1e-3,

setup_env.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+
+# Prompt the user for their name.
+echo "What is the name of the environment?"
+read -r name
+
+mkdir -p .env
+# Create a virtual environment
+python3 -m venv .env/$name
+
+# Activate the virtual environment
+source .env/$name/bin/activate
+
+# Create a temporary requirements file
+# Read the contents of the file into a variable.
+contents=$(cat requirements.txt)
+# Replace all occurrences of ">=" with "==" in the variable.
+contents=$(echo "$contents" | sed 's/>=$/==/g')
+# Write the contents of the variable to the file.
+echo "$contents" > requirements.tmp
+
+# Install the required dependencies from the temporary file
+pip install -r requirements.tmp
+
+rm requirements.tmp
+
+# Install an editable version of the package
+pip install -e .[dev]

src/pytorch_tabular/models/gate/config.py

@@ -94,10 +94,10 @@ class GatedAdditiveTreeEnsembleConfig(ModelConfig):
         default=0.0, metadata={"help": "Dropout rate for the feature abstraction layer. Defaults to 0.0"}
     )
 
-    tree_depth: int = field(default=5, metadata={"help": "Depth of the tree. Defaults to 5"})
+    tree_depth: int = field(default=4, metadata={"help": "Depth of the tree. Defaults to 5"})
 
     num_trees: int = field(
-        default=20,
+        default=10,
         metadata={"help": "Number of trees to use in the ensemble. Defaults to 20"},
     )
 

src/pytorch_tabular/models/gate/gate_model.py

@@ -95,8 +95,9 @@ def _build_network(self):
         )
         if self.tree_wise_attention:
             self.tree_attention = nn.MultiheadAttention(
-                self.output_dim,
-                1,
+                embed_dim=self.output_dim,
+                num_heads=1,
+                batch_first=False,
                 dropout=self.tree_wise_attention_dropout,
             )
 
@@ -123,7 +124,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
                 tree_input = torch.cat([tree_input, tree_output], 1)
         tree_outputs = torch.cat(tree_outputs, dim=-1)
         if self.tree_wise_attention:
-            tree_outputs, _ = self.tree_attention(tree_outputs)
+            tree_outputs = tree_outputs.permute(2, 0, 1)
+            tree_outputs, _ = self.tree_attention(tree_outputs, tree_outputs, tree_outputs)
+            tree_outputs = tree_outputs.permute(1, 2, 0)
         return tree_outputs
 
 
@@ -210,6 +213,8 @@ def _build_network(self):
             feature_mask_function=self.hparams.feature_mask_function,
             batch_norm_continuous_input=self.hparams.batch_norm_continuous_input,
             chain_trees=self.hparams.chain_trees,
+            tree_wise_attention=self.hparams.tree_wise_attention,
+            tree_wise_attention_dropout=self.hparams.tree_wise_attention_dropout,
         )
         # Embedding Layer
         self._embedding_layer = self._backbone._build_embedding_layer()