feat(model): add evaluate and refactor predict (#59)

* feat(model): add `evaluate` and refactor `predict`

* doc: update CHANGELOG.rst

* test(predict): add unit tests on `predict`

* feat: add no_grad decorator for `evaluate`

* style: minor fix for mse

* chore: minor improvement on mse

* fix: update example code

* doc: improve docstrings

* test: improve unit tests

Author: xuyxu

CHANGELOG.rst

@@ -4,6 +4,7 @@ Changelog
 Ver 0.1.*
 ---------
 
+* |Feature| |API| Improve the functionality of :meth:`evaluate` and :meth:`predict` | `@xuyxu <https://github.com/xuyxu>`__
 * |Feature| |API| Add :class:`FastGeometricClassifier` and :class:`FastGeometricRegressor` | `@xuyxu <https://github.com/xuyxu>`__
 * |Enhancement| Add flexible instantiation of optimizers and schedulers | `@cspsampedro <https://github.com/cspsampedro>`__
 * |Feature| |API| Add support on accepting instantiated base estimators as valid input | `@xuyxu <https://github.com/xuyxu>`__

examples/classification_cifar10_cnn.py

@@ -12,7 +12,6 @@
 from torchensemble.bagging import BaggingClassifier
 from torchensemble.gradient_boosting import GradientBoostingClassifier
 from torchensemble.snapshot_ensemble import SnapshotEnsembleClassifier
-from torchensemble.fast_geometric import FastGeometricClassifier
 
 from torchensemble.utils.logging import set_logger
 
@@ -115,7 +114,7 @@ def forward(self, x):
 
     # Evaluating
     tic = time.time()
-    testing_acc = model.predict(test_loader)
+    testing_acc = model.evaluate(test_loader)
     toc = time.time()
     evaluating_time = toc - tic
 
@@ -131,13 +130,15 @@ def forward(self, x):
     # Set the optimizer
     model.set_optimizer("Adam", lr=lr, weight_decay=weight_decay)
 
+    # Training
     tic = time.time()
     model.fit(train_loader, epochs=epochs)
     toc = time.time()
     training_time = toc - tic
 
+    # Evaluating
     tic = time.time()
-    testing_acc = model.predict(test_loader)
+    testing_acc = model.evaluate(test_loader)
     toc = time.time()
     evaluating_time = toc - tic
 
@@ -153,13 +154,15 @@ def forward(self, x):
     # Set the optimizer
     model.set_optimizer("Adam", lr=lr, weight_decay=weight_decay)
 
+    # Training
     tic = time.time()
     model.fit(train_loader, epochs=epochs)
     toc = time.time()
     training_time = toc - tic
 
+    # Evaluating
     tic = time.time()
-    testing_acc = model.predict(test_loader)
+    testing_acc = model.evaluate(test_loader)
     toc = time.time()
     evaluating_time = toc - tic
 
@@ -175,13 +178,15 @@ def forward(self, x):
     # Set the optimizer
     model.set_optimizer("Adam", lr=lr, weight_decay=weight_decay)
 
+    # Training
     tic = time.time()
     model.fit(train_loader, epochs=epochs)
     toc = time.time()
     training_time = toc - tic
 
+    # Evaluating
     tic = time.time()
-    testing_acc = model.predict(test_loader)
+    testing_acc = model.evaluate(test_loader)
     toc = time.time()
     evaluating_time = toc - tic
 
@@ -202,13 +207,15 @@ def forward(self, x):
     # Set the optimizer
     model.set_optimizer("Adam", lr=lr, weight_decay=weight_decay)
 
+    # Training
     tic = time.time()
     model.fit(train_loader, epochs=epochs)
     toc = time.time()
     training_time = toc - tic
 
+    # Evaluating
     tic = time.time()
-    testing_acc = model.predict(test_loader)
+    testing_acc = model.evaluate(test_loader)
     toc = time.time()
     evaluating_time = toc - tic
 

examples/fast_geometric_ensemble_cifar10_resnet18.py

@@ -168,5 +168,5 @@ def forward(self, x):
     )
 
     # Evaluate
-    acc = model.predict(test_loader)
+    acc = model.evaluate(test_loader)
     print("Testing Acc: {:.3f}".format(acc))

examples/regression_YearPredictionMSD_mlp.py

@@ -114,7 +114,7 @@ def forward(self, x):
     training_time = toc - tic
 
     tic = time.time()
-    testing_mse = model.predict(test_loader)
+    testing_mse = model.evaluate(test_loader)
     toc = time.time()
     evaluating_time = toc - tic
 
@@ -136,7 +136,7 @@ def forward(self, x):
     training_time = toc - tic
 
     tic = time.time()
-    testing_mse = model.predict(test_loader)
+    testing_mse = model.evaluate(test_loader)
     toc = time.time()
     evaluating_time = toc - tic
 
@@ -158,7 +158,7 @@ def forward(self, x):
     training_time = toc - tic
 
     tic = time.time()
-    testing_mse = model.predict(test_loader)
+    testing_mse = model.evaluate(test_loader)
     toc = time.time()
     evaluating_time = toc - tic
 
@@ -180,7 +180,7 @@ def forward(self, x):
     training_time = toc - tic
 
     tic = time.time()
-    testing_mse = model.predict(test_loader)
+    testing_mse = model.evaluate(test_loader)
     toc = time.time()
     evaluating_time = toc - tic
 
@@ -207,7 +207,7 @@ def forward(self, x):
     training_time = toc - tic
 
     tic = time.time()
-    testing_acc = model.predict(test_loader)
+    testing_acc = model.evaluate(test_loader)
     toc = time.time()
     evaluating_time = toc - tic
 

torchensemble/_base.py

@@ -3,12 +3,13 @@
 import torch
 import logging
 import warnings
+import numpy as np
 import torch.nn as nn
 
 from . import _constants as const
 
 
-def torchensemble_model_doc(header, item):
+def torchensemble_model_doc(header="", item="model"):
     """
     A decorator on obtaining documentation for different methods in the
     ensemble. This decorator is modified from `sklearn.py` in XGBoost.
@@ -27,12 +28,13 @@ def get_doc(item):
             "model": const.__model_doc,
             "seq_model": const.__seq_model_doc,
             "fit": const.__fit_doc,
+            "predict": const.__predict_doc,
             "set_optimizer": const.__set_optimizer_doc,
             "set_scheduler": const.__set_scheduler_doc,
             "classifier_forward": const.__classification_forward_doc,
-            "classifier_predict": const.__classification_predict_doc,
+            "classifier_evaluate": const.__classification_evaluate_doc,
             "regressor_forward": const.__regression_forward_doc,
-            "regressor_predict": const.__regression_predict_doc,
+            "regressor_evaluate": const.__regression_evaluate_doc,
         }
         return __doc[item]
 
@@ -45,7 +47,7 @@ def adddoc(cls):
     return adddoc
 
 
-class BaseModule(abc.ABC, nn.Module):
+class BaseModule(nn.Module):
     """Base class for all ensembles.
 
     WARNING: This class cannot be used directly.
@@ -160,13 +162,13 @@ def _validate_parameters(self, epochs, log_interval):
     @abc.abstractmethod
     def set_optimizer(self, optimizer_name, **kwargs):
         """
-        Implementation on the process of setting the optimizer.
+        Implementation on setting the parameter optimizer.
         """
 
     @abc.abstractmethod
     def set_scheduler(self, scheduler_name, **kwargs):
         """
-        Implementation on the process of setting the scheduler.
+        Implementation on setting the learning rate scheduler.
         """
 
     @abc.abstractmethod
@@ -191,8 +193,81 @@ def fit(
         Implementation on the training stage of the ensemble.
         """
 
-    @abc.abstractmethod
-    def predict(self, test_loader):
-        """
-        Implementation on the evaluating stage of the ensemble.
-        """
+    @torch.no_grad()
+    def predict(self, X, return_numpy=True):
+        """Docstrings decorated by downstream models."""
+        self.eval()
+        pred = None
+
+        if isinstance(X, torch.Tensor):
+            pred = self.forward(X.to(self.device))
+        elif isinstance(X, np.ndarray):
+            X = torch.Tensor(X).to(self.device)
+            pred = self.forward(X)
+        else:
+            msg = (
+                "The type of input X should be one of {{torch.Tensor,"
+                " np.ndarray}}."
+            )
+            raise ValueError(msg)
+
+        pred = pred.cpu()
+        if return_numpy:
+            return pred.numpy()
+
+        return pred
+
+
+class BaseClassifier(BaseModule):
+    """Base class for all ensemble classifiers.
+
+    WARNING: This class cannot be used directly.
+    Please use the derived classes instead.
+    """
+
+    @torch.no_grad()
+    def evaluate(self, test_loader, return_loss=False):
+        """Docstrings decorated by downstream models."""
+        self.eval()
+        correct = 0
+        total = 0
+        criterion = nn.CrossEntropyLoss()
+        loss = 0.0
+
+        for _, (data, target) in enumerate(test_loader):
+            data, target = data.to(self.device), target.to(self.device)
+            output = self.forward(data)
+            _, predicted = torch.max(output.data, 1)
+            correct += (predicted == target).sum().item()
+            total += target.size(0)
+            loss += criterion(output, target)
+
+        acc = 100 * correct / total
+        loss /= len(test_loader)
+
+        if return_loss:
+            return acc, float(loss)
+
+        return acc
+
+
+class BaseRegressor(BaseModule):
+    """Base class for all ensemble regressors.
+
+    WARNING: This class cannot be used directly.
+    Please use the derived classes instead.
+    """
+
+    @torch.no_grad()
+    def evaluate(self, test_loader):
+        """Docstrings decorated by downstream models."""
+        self.eval()
+        mse = 0.0
+        criterion = nn.MSELoss()
+
+        for _, (data, target) in enumerate(test_loader):
+            data, target = data.to(self.device), target.to(self.device)
+            output = self.forward(data)
+            mse += criterion(output, target)
+
+        return float(mse) / len(test_loader)

torchensemble/_constants.py

@@ -119,6 +119,27 @@
 """
 
 
+__predict_doc = """
+    Return the predictions of the ensemble given the testing data.
+
+    Parameters
+    ----------
+    X : {Tensor, ndarray}
+        A data batch in the form of tensor or Numpy array.
+    return_numpy : bool, default=True
+        Whether to convert the predictions into a Numpy array.
+
+    Returns
+    -------
+    pred : Array of shape (n_samples, n_outputs)
+        For classifiers, ``n_outputs`` is the number of distinct classes. For
+        regressors, ``n_output`` is the number of target variables.
+
+        - If ``return_numpy`` is ``False``, the result is a tensor.
+        - If ``return_numpy`` is ``True``, the result is a Numpy array.
+"""
+
+
 __classification_forward_doc = """
     Parameters
     ----------
@@ -133,17 +154,25 @@
 """
 
 
-__classification_predict_doc = """
+__classification_evaluate_doc = """
+    Compute the classification accuracy of the ensemble given the testing
+    dataloader and optionally the average cross-entropy loss.
+
     Parameters
     ----------
     test_loader : torch.utils.data.DataLoader
-        A :mod:`torch.utils.data.DataLoader` container that contains the
-        evaluating data.
+        A data loader that contains the testing data.
+    return_loss : bool, default=False
+        Whether to return the average cross-entropy loss over all batches
+        in the ``test_loader``.
 
     Returns
     -------
     accuracy : float
-        The testing accuracy of the fitted ensemble on ``test_loader``.
+        The classification accuracy of the fitted ensemble on ``test_loader``.
+    loss : float
+        The average cross-entropy loss of the fitted ensemble on
+        ``test_loader``, only available when ``return_loss`` is True.
 """
 
 
@@ -161,16 +190,18 @@
 """
 
 
-__regression_predict_doc = """
+__regression_evaluate_doc = """
+    Compute the mean squared error (MSE) of the ensemble given the testing
+    dataloader.
+
     Parameters
     ----------
     test_loader : torch.utils.data.DataLoader
-        A :mod:`torch.utils.data.DataLoader` container that contains the
-        evaluating data.
+        A data loader that contains the testing data.
 
     Returns
     -------
     mse : float
-        The testing mean squared error (MSE) of the fitted ensemble on
+        The testing mean squared error of the fitted ensemble on
         ``test_loader``.
 """