From cd483581d26eb26d56e9d2494d7e67c86d866276 Mon Sep 17 00:00:00 2001 From: SHIVANSH-ux-ys Date: Sat, 2 May 2026 01:41:19 +0530 Subject: [PATCH 1/2] feat: handle lapsed tabnet dependency with warnings and migration guide --- README.md | 2 +- docs/migration_guides/tabnet_migration.md | 34 +++++++++++++++++++ .../models/tabnet/tabnet_model.py | 8 +++++ 3 files changed, 43 insertions(+), 1 deletion(-) create mode 100644 docs/migration_guides/tabnet_migration.md diff --git a/README.md b/README.md index 2a1a5289..6b402f87 100644 --- a/README.md +++ b/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. diff --git a/docs/migration_guides/tabnet_migration.md b/docs/migration_guides/tabnet_migration.md new file mode 100644 index 00000000..e0299f82 --- /dev/null +++ b/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). diff --git a/src/pytorch_tabular/models/tabnet/tabnet_model.py b/src/pytorch_tabular/models/tabnet/tabnet_model.py index 49224048..5ed067d8 100644 --- a/src/pytorch_tabular/models/tabnet/tabnet_model.py +++ b/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): From d60afc16531c9b9edca9e03027a3bf33a551a6ce Mon Sep 17 00:00:00 2001 From: SHIVANSH-ux-ys Date: Wed, 6 May 2026 10:11:25 +0530 Subject: [PATCH 2/2] feat: implement comprehensive ONNX export validation and fix dynamic_axes bug --- pyproject.toml | 3 + src/pytorch_tabular/tabular_model.py | 62 +++++++++++++++------ tests/test_common.py | 82 +++++++++++++++++++++++++++- 3 files changed, 128 insertions(+), 19 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 3dd8a037..e0f1c7d2 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/pytorch_tabular/tabular_model.py b/src/pytorch_tabular/tabular_model.py index 5c402d1a..64bdd4a5 100644 --- a/src/pytorch_tabular/tabular_model.py +++ b/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") diff --git a/tests/test_common.py b/tests/test_common.py index 72b0c0da..154421b7 100644 --- a/tests/test_common.py +++ b/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"]])