feat: add vectorize method

Author: xuyxu

CHANGELOG.rst

@@ -18,6 +18,7 @@ Changelog
 Ver 0.1.*
 ---------
 
+* |Feature| |API| Add :meth:`vectorize` for faster ensemble model inference using :mod:`functorch` (requiring :mod:`torch` version >= 1.13.0) | `@xuyxu <https://github.com/xuyxu>`__
 * |Feature| |API| Add ``voting_strategy`` parameter for :class:`VotingClassifer`, :class:`NeuralForestClassifier`, and :class:`SnapshotEnsembleClassifier` | `@LukasGardberg <https://github.com/LukasGardberg>`__
 * |Fix| Fix the sampling issue in :class:`BaggingClassifier` and :class:`BaggingRegressor` | `@SunHaozhe <https://github.com/SunHaozhe>`__
 * |Feature| |API| Add :class:`NeuralForestClassifier` and :class:`NeuralForestRegressor` | `@xuyxu <https://github.com/xuyxu>`__

torchensemble/_base.py

@@ -29,9 +29,10 @@ def get_doc(item):
         __doc = {
             "model": const.__model_doc,
             "seq_model": const.__seq_model_doc,
-            "tree_ensmeble_model": const.__tree_ensemble_doc,
+            "tree_ensemble_model": const.__tree_ensemble_doc,
             "fit": const.__fit_doc,
             "predict": const.__predict_doc,
+            "vectorize": const.__vectorize_doc,
             "set_optimizer": const.__set_optimizer_doc,
             "set_scheduler": const.__set_scheduler_doc,
             "set_criterion": const.__set_criterion_doc,
@@ -198,6 +199,21 @@ def predict(self, *x):
         pred = self.forward(*x_device)
         pred = pred.cpu()
         return pred
+    
+    def vectorize(self):
+        """Docstrings decorated by downstream ensembles."""
+        try:
+            from functorch import combine_state_for_ensemble
+        except Exception:
+            msg = (
+                "Failed to import functorch utils, please make sure the"
+                " Pytorch version >= 1.13.0."
+            )
+            raise RuntimeError(msg)
+
+        self.eval()
+        fmodel, params, buffers = combine_state_for_ensemble(self.estimators_)
+        return fmodel, params, buffers
 
 
 class BaseTreeEnsemble(BaseModule):

torchensemble/_constants.py

@@ -173,6 +173,20 @@
 """
 
 
+__vectorize_doc = """
+    Return the vectorization result of the ensemble using functorch.
+
+    Returns
+    -------
+    fmodel : FunctionalModuleWithBuffers
+        Functional version of one of the models in the ensemble.
+    params : tuple
+        Tuple of stacked model parameters in the ensemble.
+    buffers : tuple
+        Tuple of buffers, empty if not exists.
+"""
+
+
 __classification_forward_doc = """
     Parameters
     ----------

torchensemble/voting.py

@@ -309,7 +309,7 @@ def predict(self, *x):
 
 
 @torchensemble_model_doc(
-    """Implementation on the NeuralForestClassifier.""", "tree_ensmeble_model"
+    """Implementation on the NeuralForestClassifier.""", "tree_ensemble_model"
 )
 class NeuralForestClassifier(BaseTreeEnsemble, VotingClassifier):
     def __init__(self, voting_strategy="soft", **kwargs):
@@ -561,7 +561,7 @@ def predict(self, *x):
 
 
 @torchensemble_model_doc(
-    """Implementation on the NeuralForestRegressor.""", "tree_ensmeble_model"
+    """Implementation on the NeuralForestRegressor.""", "tree_ensemble_model"
 )
 class NeuralForestRegressor(BaseTreeEnsemble, VotingRegressor):
     @torchensemble_model_doc(