feat: add internal unsqueeze operation in forward of all classifiers (#136)

Author: xuyxu

CHANGELOG.rst

@@ -18,6 +18,7 @@ Changelog
 Ver 0.1.*
 ---------
 
+* |Feature| Add internal :meth:`unsqueeze` operation in :meth:`forward` of all classifiers | `@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,7 +29,7 @@ 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,
             "set_optimizer": const.__set_optimizer_doc,

torchensemble/adversarial_training.py

@@ -226,7 +226,8 @@ class AdversarialTrainingClassifier(_BaseAdversarialTraining, BaseClassifier):
     def forward(self, *x):
         # Take the average over class distributions from all base estimators.
         outputs = [
-            F.softmax(estimator(*x), dim=1) for estimator in self.estimators_
+            F.softmax(op.unsqueeze_tensor(estimator(*x)), dim=1)
+            for estimator in self.estimators_
         ]
         proba = op.average(outputs)
 

torchensemble/bagging.py

@@ -94,7 +94,8 @@ class BaggingClassifier(BaseClassifier):
     def forward(self, *x):
         # Average over class distributions from all base estimators.
         outputs = [
-            F.softmax(estimator(*x), dim=1) for estimator in self.estimators_
+            F.softmax(op.unsqueeze_tensor(estimator(*x)), dim=1)
+            for estimator in self.estimators_
         ]
         proba = op.average(outputs)
 

torchensemble/fast_geometric.py

@@ -176,7 +176,7 @@ class FastGeometricClassifier(_BaseFastGeometric, BaseClassifier):
         "classifier_forward",
     )
     def forward(self, *x):
-        proba = self._forward(*x)
+        proba = op.unsqueeze_tensor(self._forward(*x))
 
         return F.softmax(proba, dim=1)
 

torchensemble/fusion.py

@@ -39,7 +39,7 @@ def _forward(self, *x):
         "classifier_forward",
     )
     def forward(self, *x):
-        output = self._forward(*x)
+        output = op.unsqueeze_tensor(self._forward(*x))
         proba = F.softmax(output, dim=1)
 
         return proba

torchensemble/gradient_boosting.py

@@ -420,7 +420,10 @@ def fit(
         "classifier_forward",
     )
     def forward(self, *x):
-        output = [estimator(*x) for estimator in self.estimators_]
+        output = [
+            op.unsqueeze_tensor(estimator(*x))
+            for estimator in self.estimators_
+        ]
         output = op.sum_with_multiplicative(output, self.shrinkage_rate)
         proba = F.softmax(output, dim=1)
 

torchensemble/snapshot_ensemble.py

@@ -212,6 +212,16 @@ class SnapshotEnsembleClassifier(_BaseSnapshotEnsemble, BaseClassifier):
     def __init__(self, voting_strategy="soft", **kwargs):
         super().__init__(**kwargs)
 
+        implemented_strategies = {"soft", "hard"}
+        if voting_strategy not in implemented_strategies:
+            msg = (
+                "Voting strategy {} is not implemented, "
+                "please choose from {}."
+            )
+            raise ValueError(
+                msg.format(voting_strategy, implemented_strategies)
+            )
+
         self.voting_strategy = voting_strategy
 
     @torchensemble_model_doc(
@@ -221,13 +231,13 @@ def __init__(self, voting_strategy="soft", **kwargs):
     def forward(self, *x):
 
         outputs = [
-            F.softmax(estimator(*x), dim=1) for estimator in self.estimators_
+            F.softmax(op.unsqueeze_tensor(estimator(*x)), dim=1)
+            for estimator in self.estimators_
         ]
 
         if self.voting_strategy == "soft":
             proba = op.average(outputs)
-
-        elif self.voting_strategy == "hard":
+        else:
             proba = op.majority_vote(outputs)
 
         return proba

torchensemble/soft_gradient_boosting.py

@@ -406,7 +406,10 @@ def fit(
         "classifier_forward",
     )
     def forward(self, *x):
-        output = [estimator(*x) for estimator in self.estimators_]
+        output = [
+            op.unsqueeze_tensor(estimator(*x))
+            for estimator in self.estimators_
+        ]
         output = op.sum_with_multiplicative(output, self.shrinkage_rate)
         proba = F.softmax(output, dim=1)
 

torchensemble/tests/test_all_models.py

@@ -49,7 +49,7 @@ def __init__(self):
         self.linear2 = nn.Linear(2, 2)
 
     def forward(self, X):
-        X = X.view(X.size()[0], -1)
+        X = X.view(X.size(0), -1)
         output = self.linear1(X)
         output = self.linear2(output)
         return output