initial attempt to implementing hard majority voting

torchensemble/_base.py

@@ -280,8 +280,16 @@ def evaluate(self, test_loader, return_loss=False):
 
         for _, elem in enumerate(test_loader):
             data, target = split_data_target(elem, self.device)
+
+            # Get averaged probabilities over all base estimators
+            # for each sample in batch.
+            # Output shape: (batch_size, n_classes)
             output = self.forward(*data)
+
+            # get val, idx for the class w max log-prob for each sample
+            # predicted shape: (batch_size)
             _, predicted = torch.max(output.data, 1)
+
             correct += (predicted == target).sum().item()
             total += target.size(0)
             loss += self._criterion(output, target)
@@ -349,8 +357,8 @@ def __init__(self, input_dim, output_dim, depth=5, lamda=1e-3, cuda=False):
 
         self._validate_parameters()
 
-        self.internal_node_num_ = 2 ** self.depth - 1
-        self.leaf_node_num_ = 2 ** self.depth
+        self.internal_node_num_ = 2**self.depth - 1
+        self.leaf_node_num_ = 2**self.depth
 
         # Different penalty coefficients for nodes in different layers
         self.penalty_list = [
@@ -420,7 +428,7 @@ def _cal_penalty(self, layer_idx, _mu, _path_prob):
         penalty = torch.tensor(0.0).to(self.device)
 
         batch_size = _mu.size()[0]
-        _mu = _mu.view(batch_size, 2 ** layer_idx)
+        _mu = _mu.view(batch_size, 2**layer_idx)
         _path_prob = _path_prob.view(batch_size, 2 ** (layer_idx + 1))
 
         for node in range(0, 2 ** (layer_idx + 1)):

torchensemble/voting.py

@@ -92,17 +92,52 @@ 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)
+
+        self._implemented_voting_strategies = ["soft", "hard"]
+
+        if voting_strategy not in self._implemented_voting_strategies:
+            msg = (
+                "Voting strategy {} is not implemented, "
+                "please choose from {}."
+            )
+            raise ValueError(
+                msg.format(
+                    voting_strategy, self._implemented_voting_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.
+
+        # output: (n_estimators, batch_size, n_classes)
         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":
+            # Do hard majority voting
+            # votes: (batch_size)
+            votes = torch.stack(outputs).argmax(dim=2).mode(dim=0)[0]
+            # Set the probability of the most voted class to 1
+            proba = torch.zeros_like(outputs[0])
+            proba.scatter_(1, votes.view(-1, 1), 1)
+
+            # TODO: make sure this is compatible with
+            # predict and evaluate
+
+        # Returns averaged class probabilities for each sample
+        # proba shape: (batch_size, n_classes)
         return proba
 
     @torchensemble_model_doc(
@@ -167,12 +202,20 @@ def fit(
         # Utils
         best_acc = 0.0
 
-        # Internal helper function on pesudo forward
+        # Internal helper function on psuedo 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)
+
+            # Maybe the voting strategy can be packaged as a function?
+            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 +452,7 @@ def fit(
         # Utils
         best_loss = float("inf")
 
-        # Internal helper function on pesudo forward
+        # Internal helper function on psuedo forward
         def _forward(estimators, *x):
             outputs = [estimator(*x) for estimator in estimators]
             pred = op.average(outputs)