initial attempt to implementing hard majority voting

torchensemble/_base.py

@@ -280,8 +280,11 @@ def evaluate(self, test_loader, return_loss=False):
 
         for _, elem in enumerate(test_loader):
             data, target = split_data_target(elem, self.device)
+
             output = self.forward(*data)
+
             _, predicted = torch.max(output.data, 1)
+
             correct += (predicted == target).sum().item()
             total += target.size(0)
             loss += self._criterion(output, target)

torchensemble/voting.py

@@ -92,16 +92,39 @@ def _parallel_fit_per_epoch(
     """Implementation on the VotingClassifier.""", "model"
 )
 class VotingClassifier(BaseClassifier):
+    def __init__(self, voting_strategy="soft", **kwargs):
+        super(VotingClassifier, self).__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(
         """Implementation on the data forwarding in VotingClassifier.""",
         "classifier_forward",
     )
     def forward(self, *x):
         # Average over class distributions from all base estimators.
+
         outputs = [
             F.softmax(estimator(*x), dim=1) for estimator in self.estimators_
         ]
-        proba = op.average(outputs)
+
+        if self.voting_strategy == "soft":
+            proba = op.average(outputs)
+
+        elif self.voting_strategy == "hard":
+            votes = torch.stack(outputs).argmax(dim=2).mode(dim=0)[0]
+            proba = torch.zeros_like(outputs[0])
+            proba.scatter_(1, votes.view(-1, 1), 1)
 
         return proba
 
@@ -167,12 +190,19 @@ def fit(
         # Utils
         best_acc = 0.0
 
-        # Internal helper function on pesudo forward
+        # Internal helper function on pseudo forward
         def _forward(estimators, *x):
             outputs = [
                 F.softmax(estimator(*x), dim=1) for estimator in estimators
             ]
-            proba = op.average(outputs)
+
+            if self.voting_strategy == "soft":
+                proba = op.average(outputs)
+
+            elif self.voting_strategy == "hard":
+                votes = torch.stack(outputs).argmax(dim=2).mode(dim=0)[0]
+                proba = torch.zeros_like(outputs[0])
+                proba.scatter_(1, votes.view(-1, 1), 1)
 
             return proba
 
@@ -409,7 +439,7 @@ def fit(
         # Utils
         best_loss = float("inf")
 
-        # Internal helper function on pesudo forward
+        # Internal helper function on pseudo forward
         def _forward(estimators, *x):
             outputs = [estimator(*x) for estimator in estimators]
             pred = op.average(outputs)