-- added feature importance to FT Transformer

Author: manujosephv

examples/to_test_classification.py

@@ -99,25 +99,25 @@
 #     metrics=["f1", "accuracy"],
 #     metrics_params=[{"num_classes": num_classes, "average": "macro"}, {}],
 # )
-# model_config = TabTransformerConfig(
+model_config = TabTransformerConfig(
+    task="classification",
+    metrics=["f1", "accuracy"],
+    share_embedding = True,
+    share_embedding_strategy="add",
+    shared_embedding_fraction=0.25,
+    metrics_params=[{"num_classes": num_classes, "average": "macro"}, {}],
+)
+# model_config = FTTransformerConfig(
 #     task="classification",
 #     metrics=["f1", "accuracy"],
+#     # embedding_initialization=None,
+#     embedding_bias=True,
 #     share_embedding = True,
 #     share_embedding_strategy="fraction",
 #     shared_embedding_fraction=0.25,
 #     metrics_params=[{"num_classes": num_classes, "average": "macro"}, {}],
 # )
-model_config = FTTransformerConfig(
-    task="classification",
-    metrics=["f1", "accuracy"],
-    # embedding_initialization=None,
-    embedding_bias=False,
-    share_embedding = True,
-    share_embedding_strategy="fraction",
-    shared_embedding_fraction=0.25,
-    metrics_params=[{"num_classes": num_classes, "average": "macro"}, {}],
-)
-trainer_config = TrainerConfig(gpus=-1, auto_select_gpus=True, fast_dev_run=False, max_epochs=5, batch_size=512)
+trainer_config = TrainerConfig(gpus=-1, auto_select_gpus=True, fast_dev_run=True, max_epochs=5, batch_size=512)
 experiment_config = ExperimentConfig(project_name="PyTorch Tabular Example", 
                                      run_name="node_forest_cov", 
                                      exp_watch="gradients", 
@@ -147,9 +147,14 @@
     # loss=cust_loss,
     train_sampler=sampler)
 
-result = tabular_model.evaluate(test)
-print(result)
-# test.drop(columns=target_name, inplace=True)
+from pytorch_tabular.categorical_encoders import CategoricalEmbeddingTransformer
+transformer = CategoricalEmbeddingTransformer(tabular_model)
+train_transform = transformer.fit_transform(train)
+# test_transform = transformer.transform(test)
+# ft = tabular_model.model.feature_importance()
+# result = tabular_model.evaluate(test)
+# print(result)
+# test.drop(columns=ta6rget_name, inplace=True)
 # pred_df = tabular_model.predict(test)
 # print(pred_df.head())
 # pred_df.to_csv("output/temp2.csv")

pytorch_tabular/models/base_model.py

@@ -121,7 +121,7 @@ def calculate_metrics(self, y, y_hat, tag):
         for metric, metric_str, metric_params in zip(
             self.metrics, self.hparams.metrics, self.hparams.metrics_params
         ):
-            if (self.hparams.task == "regression") and (self.hparams.output_dim > 1):
+            if (self.hparams.task == "regression"):
                 _metrics = []
                 for i in range(self.hparams.output_dim):
                     if (

pytorch_tabular/models/common.py

@@ -150,7 +150,7 @@ class MultiHeadedAttention(nn.Module):
     Multi Headed Attention Block in Transformers
     """
     def __init__(
-        self, input_dim: int, num_heads: int = 8, head_dim: int = 16, dropout: int = 0.1
+        self, input_dim: int, num_heads: int = 8, head_dim: int = 16, dropout: int = 0.1, keep_attn: bool = True
     ):
         super().__init__()
         assert (
@@ -159,6 +159,7 @@ def __init__(
         inner_dim = head_dim * num_heads
         self.n_heads = num_heads
         self.scale = head_dim ** -0.5
+        self.keep_attn = keep_attn
 
         self.to_qkv = nn.Linear(input_dim, inner_dim * 3, bias=False)
         self.to_out = nn.Linear(inner_dim, input_dim)
@@ -173,7 +174,8 @@ def forward(self, x):
 
         attn = sim.softmax(dim=-1)
         attn = self.dropout(attn)
-
+        if self.keep_attn:
+            self.attn_weights = attn
         out = einsum("b h i j, b h j d -> b h i d", attn, v)
         out = rearrange(out, "b h n d -> b n (h d)", h=h)
         return self.to_out(out)
@@ -211,7 +213,15 @@ def forward(self, X: torch.Tensor) -> torch.Tensor:
         else:
             out[:, : shared_embed.shape[1]] = shared_embed
         return out
-
+    
+    @property
+    def weight(self):
+        w = self.embed.weight.detach()
+        if self.add_shared_embed:
+            w += self.shared_embed
+        else:
+            w[:, : self.shared_embed.shape[1]] = self.shared_embed
+        return w
 
 class TransformerEncoderBlock(nn.Module):
     """A single Transformer Encoder Block
@@ -223,6 +233,7 @@ def __init__(
         ff_hidden_multiplier: int = 4,
         ff_activation: str = "GEGLU",
         attn_dropout: float = 0.1,
+        keep_attn: bool = True,
         ff_dropout: float = 0.1,
         add_norm_dropout: float = 0.1,
         transformer_head_dim: Optional[int] = None,
@@ -235,6 +246,7 @@ def __init__(
             if transformer_head_dim is None
             else transformer_head_dim,
             dropout=attn_dropout,
+            keep_attn = keep_attn
         )
 
         try:

pytorch_tabular/models/ft_transformer/config.py

@@ -128,6 +128,12 @@ class FTTransformerConfig(ModelConfig):
             "help": "Fraction of the input_embed_dim to be reserved by the shared embedding. Should be less than one. Defaults to 0.25"
         },
     )
+    attn_feature_importance: bool = field(
+        default = True,
+        metadata={
+            "help": "If you are facing memory issues, you can turn off feature importance which will not save the attention weights. Defaults to True"
+        },
+    )
     num_heads: int = field(
         default=8,
         metadata={

pytorch_tabular/models/ft_transformer/ft_transformer.py

@@ -3,10 +3,11 @@
 # For license information, see LICENSE.TXT
 """Feature Tokenizer Transformer Model"""
 import logging
-from collections import OrderedDict
 import math
+from collections import OrderedDict
 from typing import Dict
 
+import pandas as pd
 import pytorch_lightning as pl
 import torch
 import torch.nn as nn
@@ -118,9 +119,11 @@ def _build_network(self):
                 attn_dropout=self.hparams.attn_dropout,
                 ff_dropout=self.hparams.ff_dropout,
                 add_norm_dropout=self.hparams.add_norm_dropout,
+                keep_attn=self.hparams.attn_feature_importance #Can use Attn Weights to derive feature importance
             )
         self.transformer_blocks = nn.Sequential(self.transformer_blocks)
-        self.attention_weights = [None] * self.hparams.num_attn_blocks
+        if self.hparams.attn_feature_importance:
+            self.attention_weights_ = [None] * self.hparams.num_attn_blocks
         if self.hparams.batch_norm_continuous_input:
             self.normalizing_batch_norm = nn.BatchNorm1d(self.hparams.continuous_dim)
         # Final MLP Layers
@@ -177,11 +180,31 @@ def forward(self, x: Dict):
         x = self.add_cls(x)
         for i, block in enumerate(self.transformer_blocks):
             x = block(x)
+            if self.hparams.attn_feature_importance:
+                self.attention_weights_[i] = block.mha.attn_weights
+                # self.feature_importance_+=block.mha.attn_weights[:,:,:,-1].sum(dim=1)
+                # self._calculate_feature_importance(block.mha.attn_weights)
+        if self.hparams.attn_feature_importance:
+            self._calculate_feature_importance()
         # Flatten (Batch, N_Categorical, Hidden) --> (Batch, N_CategoricalxHidden)
         # x = rearrange(x, "b n h -> b (n h)")
         # Taking only CLS token for the prediction head
         x = self.linear_layers(x[:, -1])
         return x
+    
+    #Not Tested Properly
+    def _calculate_feature_importance(self):
+        # if self.feature_importance_.device != self.device:
+        #     self.feature_importance_ = self.feature_importance_.to(self.device)
+
+        n, h, f, _ = self.attention_weights_[0].shape
+        L = len(self.attention_weights_)
+        self.local_feature_importance = torch.zeros((n,f), device=self.device)
+        for attn_weights in self.attention_weights_:
+            self.local_feature_importance+=attn_weights[:,:,:,-1].sum(dim=1)
+        self.local_feature_importance = (1/(h*L))*self.local_feature_importance[:,:-1]
+        self.feature_importance_ = self.local_feature_importance.mean(dim=0)
+        # self.feature_importance_count_+=attn_weights.shape[0]
 
 
 class FTTransformerModel(BaseModel):
@@ -221,3 +244,10 @@ def extract_embedding(self):
             raise ValueError(
                 "Model has been trained with no categorical feature and therefore can't be used as a Categorical Encoder"
             )
+    
+    def feature_importance(self):
+        if self.hparams.attn_feature_importance:
+            importance_df = pd.DataFrame({"Features": self.hparams.categorical_cols+self.hparams.continuous_cols, "importance": self.backbone.feature_importance_.detach().cpu().numpy()})
+            return importance_df
+        else:
+            raise ValueError("If you want Feature Importance, `attn_feature_weights` should be `True`.")

pytorch_tabular/models/tab_transformer/tab_transformer.py

@@ -75,6 +75,7 @@ def _build_network(self):
                 attn_dropout=self.hparams.attn_dropout,
                 ff_dropout=self.hparams.ff_dropout,
                 add_norm_dropout=self.hparams.add_norm_dropout,
+                keep_attn = False # No easy way to convert TabTransformer Attn Weights to Feature Importance 
             )
         self.transformer_blocks = nn.Sequential(self.transformer_blocks)
         self.attention_weights = [None] * self.hparams.num_attn_blocks

pytorch_tabular/tabular_datamodule.py

@@ -232,14 +232,20 @@ def preprocess_data(
         # Target Transforms
         if all([col in data.columns for col in self.config.target]):
             if self.do_target_transform:
-                target_transforms = []
-                for col in self.config.target:
-                    _target_transform = copy.deepcopy(self.target_transform_template)
-                    data[col] = _target_transform.fit_transform(
-                        data[col].values.reshape(-1, 1)
-                    )
-                    target_transforms.append(_target_transform)
-                self.target_transforms = target_transforms
+                if stage == "fit":
+                    target_transforms = []
+                    for col in self.config.target:
+                        _target_transform = copy.deepcopy(self.target_transform_template)
+                        data[col] = _target_transform.fit_transform(
+                            data[col].values.reshape(-1, 1)
+                        )
+                        target_transforms.append(_target_transform)
+                    self.target_transforms = target_transforms
+                else:
+                    for col, _target_transform in zip(self.config.target, self.target_transforms):
+                        data[col] = _target_transform.transform(
+                            data[col].values.reshape(-1, 1)
+                        )
         return data, added_features
 
     def setup(self, stage: Optional[str] = None) -> None:

pytorch_tabular/tabular_model.py

@@ -184,7 +184,10 @@ def _get_run_name_uid(self) -> Tuple[str, int]:
         """
         if hasattr(self.config, "run_name") and self.config.run_name is not None:
             name = self.config.run_name
-        elif hasattr(self.config, "checkpoints_name") and self.config.checkpoints_name is not None:
+        elif (
+            hasattr(self.config, "checkpoints_name")
+            and self.config.checkpoints_name is not None
+        ):
             name = self.config.checkpoints_name
         else:
             name = self.config.task
@@ -287,7 +290,6 @@ def _prepare_model(self, loss, metrics, optimizer, optimizer_params, reset):
             )
             # Data Aware Initialization(for the models that need it)
             self.model.data_aware_initialization(self.datamodule)
-            
 
     def _prepare_trainer(self, max_epochs=None, min_epochs=None):
         logger.info("Preparing the Trainer...")
@@ -297,7 +299,7 @@ def _prepare_trainer(self, max_epochs=None, min_epochs=None):
             self.config.min_epochs = min_epochs
         # TODO get Trainer Arguments from the init signature
         trainer_sig = inspect.signature(pl.Trainer.__init__)
-        trainer_args = [p for p in trainer_sig.parameters.keys() if p!="self"]
+        trainer_args = [p for p in trainer_sig.parameters.keys() if p != "self"]
         trainer_args_config = {
             k: v for k, v in self.config.items() if k in trainer_args
         }
@@ -314,9 +316,14 @@ def load_best_model(self):
         if self.trainer.checkpoint_callback is not None:
             logger.info("Loading the best model...")
             ckpt_path = self.trainer.checkpoint_callback.best_model_path
-            logger.debug(f"Model Checkpoint: {ckpt_path}")
-            ckpt = pl_load(ckpt_path, map_location=lambda storage, loc: storage)
-            self.model.load_state_dict(ckpt["state_dict"])
+            if ckpt_path != "":
+                logger.debug(f"Model Checkpoint: {ckpt_path}")
+                ckpt = pl_load(ckpt_path, map_location=lambda storage, loc: storage)
+                self.model.load_state_dict(ckpt["state_dict"])
+            else:
+                logger.info(
+                    "No best model available to load. Did you run it more than 1 epoch?..."
+                )
         else:
             logger.info(
                 "No best model available to load. Did you run it more than 1 epoch?..."
@@ -737,19 +744,18 @@ def load_from_checkpoint(cls, dir: str):
         custom_params = joblib.load(os.path.join(dir, "custom_params.sav"))
         model_args = {}
         if custom_params.get("custom_loss") is not None:
-            model_args['loss'] = "MSELoss"
+            model_args["loss"] = "MSELoss"
         if custom_params.get("custom_metrics") is not None:
-            model_args['metrics'] = ["mean_squared_error"]
-            model_args['metric_params'] = [{}]
+            model_args["metrics"] = ["mean_squared_error"]
+            model_args["metric_params"] = [{}]
         if custom_params.get("custom_optimizer") is not None:
-            model_args['optimizer'] = "Adam"
+            model_args["optimizer"] = "Adam"
         if custom_params.get("custom_optimizer_params") is not None:
-            model_args['optimizer_params'] = {}
-        
+            model_args["optimizer_params"] = {}
+
         # Initializing with default metrics, losses, and optimizers. Will revert once initialized
         model = model_callable.load_from_checkpoint(
-            checkpoint_path=os.path.join(dir, "model.ckpt"),
-            **model_args
+            checkpoint_path=os.path.join(dir, "model.ckpt"), **model_args
         )
         # else:
         #     # Initializing with default values