feat: Implement Comprehensive ONNX Export Testing and Validation (#654)

README.md

@@ -71,7 +71,7 @@ For complete Documentation with tutorials visit [ReadTheDocs](https://pytorch-ta
 
 - FeedForward Network with Category Embedding is a simple FF network, but with an Embedding layers for the categorical columns.
 - [Neural Oblivious Decision Ensembles for Deep Learning on Tabular Data](https://arxiv.org/abs/1909.06312) is a model presented in ICLR 2020 and according to the authors have beaten well-tuned Gradient Boosting models on many datasets.
-- [TabNet: Attentive Interpretable Tabular Learning](https://arxiv.org/abs/1908.07442) is another model coming out of Google Research which uses Sparse Attention in multiple steps of decision making to model the output.
+- [TabNet: Attentive Interpretable Tabular Learning](https://arxiv.org/abs/1908.07442) is another model coming out of Google Research which uses Sparse Attention in multiple steps of decision making to model the output. **(Note: The `pytorch-tabnet` dependency has lapsed maintenance. See [Migration Guide](docs/migration_guides/tabnet_migration.md))**
 - [Mixture Density Networks](https://publications.aston.ac.uk/id/eprint/373/1/NCRG_94_004.pdf) is a regression model which uses gaussian components to approximate the target function and  provide a probabilistic prediction out of the box.
 - [AutoInt: Automatic Feature Interaction Learning via Self-Attentive Neural Networks](https://arxiv.org/abs/1810.11921) is a model which tries to learn interactions between the features in an automated way and create a better representation and then use this representation in downstream task
 - [TabTransformer](https://arxiv.org/abs/2012.06678) is an adaptation of the Transformer model for Tabular Data which creates contextual representations for categorical features.

docs/migration_guides/tabnet_migration.md

@@ -0,0 +1,34 @@
+# TabNet Migration Guide
+
+As of 2024, the `pytorch-tabnet` package, which is a soft dependency for the `TabNetModel` in PyTorch Tabular, has lapsed maintenance (last updated in 2023). While the model remains functional in current environments, we recommend users transition to more modern and actively maintained architectures to ensure long-term stability and performance.
+
+## Why Migrate?
+- **Maintenance Status**: `pytorch-tabnet` is no longer receiving security patches or performance updates.
+- **Compatibility**: Future updates to PyTorch or PyTorch Lightning may break `pytorch-tabnet` internals.
+- **Superior Alternatives**: Architectures like **GANDALF** and **FT-Transformer** often provide equal or superior performance on tabular benchmarks with better stability.
+
+## Recommended Alternatives
+
+### 1. GANDALF (Gated Adaptive Network for Deep Automated Learning of Features)
+GANDALF is highly efficient and often outperforms TabNet on modern benchmarks.
+- **When to use**: High performance, minimal hyperparameter tuning.
+- **Config**:
+```python
+from pytorch_tabular.models import GatedAdditiveTreeEnsembleConfig
+model_config = GatedAdditiveTreeEnsembleConfig(...)
+```
+
+### 2. FT-Transformer
+A robust adaptation of the Transformer architecture for tabular data.
+- **When to use**: When you need strong contextual representations of features.
+- **Config**:
+```python
+from pytorch_tabular.models import FTTransformerConfig
+model_config = FTTransformerConfig(...)
+```
+
+## Migration Steps
+1. Update your `ModelConfig` to use one of the recommended alternatives.
+2. If you have a pre-trained TabNet model, you will need to retrain a new model with the new architecture, as weights are not transferable between different model classes.
+
+For more information, please refer to our [Issue #611](https://github.com/manujosephv/pytorch_tabular/issues/611).

pyproject.toml

@@ -73,8 +73,11 @@ extra = [
     "kaleido>=0.2.0,<0.3.0",
     "captum>=0.5.0,<0.8.0",
     "pytorch-tabnet<4.2",
+    "onnx>=1.12.0",
+    "onnxruntime>=1.12.0",
 ]
 
+
 notebooks = [
     "ipywidgets",
     "matplotlib>3.1",

src/pytorch_tabular/models/tabnet/tabnet_model.py

@@ -3,6 +3,7 @@
 # For license information, see LICENSE.TXT
 """TabNet Model."""
 
+import warnings
 from typing import Dict
 
 import torch
@@ -87,6 +88,13 @@ def __init__(self, config: DictConfig, **kwargs):
         ], "TabNet is only implemented for Regression and Classification"
         super().__init__(config, **kwargs)
         _check_soft_dependencies("pytorch-tabnet", obj=self)
+        warnings.warn(
+            "pytorch-tabnet has lapsed maintenance in 2023. It is recommended to use "
+            "other models like GANDALF or FT-Transformer. See the migration guide "
+            "for more details: https://pytorch-tabular.readthedocs.io/en/latest/migration_guides/tabnet_migration/",
+            FutureWarning,
+            stacklevel=2,
+        )
 
     @property
     def backbone(self):

src/pytorch_tabular/tabular_model.py

@@ -1613,7 +1613,6 @@ def load_weights(self, path: Union[str, Path]) -> None:
         """
         self._load_weights(self.model, path)
 
-    # TODO Need to test ONNX export
     def save_model_for_inference(
         self,
         path: Union[str, Path],
@@ -1636,26 +1635,53 @@ def save_model_for_inference(
             torch.save(self.model, str(path))
             return True
         elif kind == "onnx":
+            try:
+                import onnx
+            except ImportError:
+                raise ImportError("onnx is not installed. Please install onnx using `pip install onnx`.")
             # Export the model
-            onnx_export_params["input_names"] = ["categorical", "continuous"]
+            input_names = []
+            x = {}
+            dynamic_axes = {}
+            if len(self.config.categorical_cols) > 0:
+                input_names.append("categorical")
+                x["categorical"] = torch.zeros(
+                    self.config.batch_size,
+                    len(self.config.categorical_cols),
+                    dtype=torch.long,
+                )
+                dynamic_axes["categorical"] = {0: "batch_size"}
+            if len(self.config.continuous_cols) > 0:
+                input_names.append("continuous")
+                x["continuous"] = torch.randn(
+                    self.config.batch_size,
+                    len(self.config.continuous_cols),
+                    requires_grad=True,
+                )
+                dynamic_axes["continuous"] = {0: "batch_size"}
+
+            onnx_export_params["input_names"] = onnx_export_params.get("input_names", input_names)
             onnx_export_params["output_names"] = onnx_export_params.get("output_names", ["output"])
-            onnx_export_params["dynamic_axes"] = {
-                onnx_export_params["input_names"][0]: {0: "batch_size"},
-                onnx_export_params["output_names"][0]: {0: "batch_size"},
-            }
-            cat = torch.zeros(
-                self.config.batch_size,
-                len(self.config.categorical_cols),
-                dtype=torch.int,
-            )
-            cont = torch.randn(
-                self.config.batch_size,
-                len(self.config.continuous_cols),
-                requires_grad=True,
-            )
-            x = {"continuous": cont, "categorical": cat}
-            torch.onnx.export(self.model, x, str(path), **onnx_export_params)
+
+            # Merging with user provided dynamic_axes if any
+            if "dynamic_axes" not in onnx_export_params:
+                onnx_export_params["dynamic_axes"] = dynamic_axes
+                for out in onnx_export_params["output_names"]:
+                    onnx_export_params["dynamic_axes"][out] = {0: "batch_size"}
+
+            # Set model to eval mode
+            curr_mode = self.model.training
+            self.model.eval()
+            try:
+                torch.onnx.export(self.model, x, str(path), **onnx_export_params)
+            finally:
+                self.model.train(curr_mode)
+
+            # Check model
+            onnx_model = onnx.load(str(path))
+            onnx.checker.check_model(onnx_model)
             return True
+
         else:
             raise ValueError("`kind` must be either pytorch or onnx")
 

tests/test_common.py

@@ -32,6 +32,8 @@
 
 # todo: move the logic to skip soft dependency dependent estimators to tags etc
 TABNET_AVAILABLE = _check_soft_dependencies("pytorch-tabnet", severity="none")
+ONNX_AVAILABLE = _check_soft_dependencies(["onnx", "onnxruntime"], severity="none")
+
 
 
 MODEL_CONFIG_SAVE_TEST = [
@@ -53,6 +55,8 @@
             "num_attn_blocks": 1,
         },
     ),
+    (GANDALFConfig, {}),
+    (FTTransformerConfig, {"num_heads": 1, "num_attn_blocks": 1}),
 ]
 MODEL_CONFIG_FEATURE_EXT_TEST = [
     CategoryEmbeddingModelConfig,
@@ -380,7 +384,7 @@ def test_save_load_statedict(
 @pytest.mark.parametrize("custom_metrics", [None, [fake_metric]])
 @pytest.mark.parametrize("custom_loss", [None, torch.nn.L1Loss()])
 @pytest.mark.parametrize("custom_optimizer", [None, torch.optim.Adagrad, "SGD", "torch_optimizer.AdaBound"])
-@pytest.mark.parametrize("save_type", ["pytorch"])  # "onnx"
+@pytest.mark.parametrize("save_type", ["pytorch"])
 def test_save_for_inference(
     regression_data,
     model_config_class,
@@ -434,6 +438,82 @@ def test_save_for_inference(
     assert os.path.exists(sv_dir / model_name)
 
 
+@pytest.mark.skipif(not ONNX_AVAILABLE, reason="ONNX or ONNXRuntime not installed")
+@pytest.mark.parametrize("model_config_class", MODEL_CONFIG_SAVE_ONNX_TEST)
+@pytest.mark.parametrize("continuous_cols", [list(DATASET_CONTINUOUS_COLUMNS)])
+@pytest.mark.parametrize("categorical_cols", [["HouseAgeBin"]])
+def test_save_for_inference_onnx(
+    regression_data,
+    model_config_class,
+    continuous_cols,
+    categorical_cols,
+    tmpdir,
+):
+    (train, test, target) = regression_data
+    data_config = DataConfig(
+        target=target,
+        continuous_cols=continuous_cols,
+        categorical_cols=categorical_cols,
+    )
+    model_config_class, model_config_params = model_config_class
+    model_config_params["task"] = "regression"
+    model_config = model_config_class(**model_config_params)
+    trainer_config = TrainerConfig(
+        max_epochs=1,
+        checkpoints=None,
+        early_stopping=None,
+        accelerator="cpu",
+        fast_dev_run=True,
+    )
+    optimizer_config = OptimizerConfig()
+
+    tabular_model = TabularModel(
+        data_config=data_config,
+        model_config=model_config,
+        optimizer_config=optimizer_config,
+        trainer_config=trainer_config,
+    )
+    tabular_model.fit(
+        train=train,
+    )
+    sv_dir = tmpdir.mkdir("saved_model")
+    model_name = "model.onnx"
+
+    # Test Export
+    tabular_model.save_model_for_inference(
+        sv_dir / model_name,
+        kind="onnx",
+    )
+    assert os.path.exists(sv_dir / model_name)
+
+    # Test Validation
+    import onnxruntime as ort
+
+    # Get PyTorch Prediction (Point predictions)
+    test_data = test.head(5)
+    pt_preds = tabular_model.predict(test_data)
+    target_col = tabular_model.config.target[0]
+    pt_vals = pt_preds[f"{target_col}_prediction"].values.reshape(-1, 1)
+
+    # Get ONNX Prediction
+    ort_sess = ort.InferenceSession(str(sv_dir / model_name))
+    inference_dataloader = tabular_model.datamodule.prepare_inference_dataloader(test_data)
+    batch = next(iter(inference_dataloader))
+
+    ort_inputs = {}
+    if len(tabular_model.config.categorical_cols) > 0:
+        ort_inputs["categorical"] = batch["categorical"].numpy().astype(np.int64)
+    if len(tabular_model.config.continuous_cols) > 0:
+        ort_inputs["continuous"] = batch["continuous"].numpy().astype(np.float32)
+    ort_outs = ort_sess.run(None, ort_inputs)
+    ort_vals = ort_outs[0]
+
+
+    # Compare results
+    np.testing.assert_allclose(pt_vals, ort_vals, rtol=1e-3, atol=1e-3)
+    assert ort_vals.shape == (5, 1)
+
+
 @pytest.mark.parametrize("model_config_class", MODEL_CONFIG_FEATURE_EXT_TEST)
 @pytest.mark.parametrize("continuous_cols", [list(DATASET_CONTINUOUS_COLUMNS)])
 @pytest.mark.parametrize("categorical_cols", [["HouseAgeBin"]])