chore: code cleaning up (#71)

* chore: code cleaning up

* update

* Update adversarial_training.py

Author: xuyxu

torchensemble/_base.py

@@ -95,37 +95,12 @@ def __getitem__(self, index):
         """Return the `index`-th base estimator in the ensemble."""
         return self.estimators_[index]
 
-    def _decide_n_outputs(self, train_loader, is_classification=True):
-        """
-        Decide the number of outputs according to the `train_loader`.
-
-        - If `is_classification` is True, the number of outputs equals the
-          number of distinct classes.
-        - If `is_classification` is False, the number of outputs equals the
-          number of target variables (e.g., `1` in univariate regression).
-        """
-        if is_classification:
-            if hasattr(train_loader.dataset, "classes"):
-                n_outputs = len(train_loader.dataset.classes)
-            # Infer `n_outputs` from the dataloader
-            else:
-                labels = []
-                for _, (_, target) in enumerate(train_loader):
-                    labels.append(target)
-                labels = torch.unique(torch.cat(labels))
-                n_outputs = labels.size(0)
-        else:
-            for _, (_, target) in enumerate(train_loader):
-                if len(target.size()) == 1:
-                    n_outputs = 1
-                else:
-                    n_outputs = target.size(1)
-                break
-
-        return n_outputs
+    @abc.abstractmethod
+    def _decide_n_outputs(self, train_loader):
+        """Decide the number of outputs according to the `train_loader`."""
 
     def _make_estimator(self):
-        """Make and configure a copy of the `self.base_estimator_`."""
+        """Make and configure a copy of `self.base_estimator_`."""
 
         # Call `deepcopy` to make a base estimator
         if not isinstance(self.base_estimator_, type):
@@ -161,17 +136,16 @@ def _validate_parameters(self, epochs, log_interval):
             self.logger.error(msg.format(log_interval))
             raise ValueError(msg.format(log_interval))
 
-    @abc.abstractmethod
     def set_optimizer(self, optimizer_name, **kwargs):
-        """
-        Implementation on setting the parameter optimizer.
-        """
+        """Set the parameter optimizer."""
+        self.optimizer_name = optimizer_name
+        self.optimizer_args = kwargs
 
-    @abc.abstractmethod
     def set_scheduler(self, scheduler_name, **kwargs):
-        """
-        Implementation on setting the learning rate scheduler.
-        """
+        """Set the learning rate scheduler."""
+        self.scheduler_name = scheduler_name
+        self.scheduler_args = kwargs
+        self.use_scheduler_ = True
 
     @abc.abstractmethod
     def forward(self, x):
@@ -227,6 +201,24 @@ class BaseClassifier(BaseModule):
     Please use the derived classes instead.
     """
 
+    def _decide_n_outputs(self, train_loader):
+        """
+        Decide the number of outputs according to the `train_loader`.
+        The number of outputs equals the number of distinct classes for
+        classifiers.
+        """
+        if hasattr(train_loader.dataset, "classes"):
+            n_outputs = len(train_loader.dataset.classes)
+        # Infer `n_outputs` from the dataloader
+        else:
+            labels = []
+            for _, (_, target) in enumerate(train_loader):
+                labels.append(target)
+            labels = torch.unique(torch.cat(labels))
+            n_outputs = labels.size(0)
+
+        return n_outputs
+
     @torch.no_grad()
     def evaluate(self, test_loader, return_loss=False):
         """Docstrings decorated by downstream models."""
@@ -260,6 +252,21 @@ class BaseRegressor(BaseModule):
     Please use the derived classes instead.
     """
 
+    def _decide_n_outputs(self, train_loader):
+        """
+        Decide the number of outputs according to the `train_loader`.
+        The number of outputs equals the number of target variables for
+        regressors (e.g., `1` in univariate regression).
+        """
+        for _, (_, target) in enumerate(train_loader):
+            if len(target.size()) == 1:
+                n_outputs = 1
+            else:
+                n_outputs = target.size(1)
+            break
+
+        return n_outputs
+
     @torch.no_grad()
     def evaluate(self, test_loader):
         """Docstrings decorated by downstream models."""

torchensemble/adversarial_training.py

@@ -23,11 +23,7 @@
 from .utils import operator as op
 
 
-__all__ = [
-    "_BaseAdversarialTraining",
-    "AdversarialTrainingClassifier",
-    "AdversarialTrainingRegressor",
-]
+__all__ = ["AdversarialTrainingClassifier", "AdversarialTrainingRegressor"]
 
 
 __fit_doc = """
@@ -218,10 +214,6 @@ def _validate_parameters(self, epochs, epsilon, log_interval):
     "model",
 )
 class AdversarialTrainingClassifier(_BaseAdversarialTraining, BaseClassifier):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-        self.is_classification = True
-
     @torchensemble_model_doc(
         """Implementation on the data forwarding in AdversarialTrainingClassifier.""",  # noqa: E501
         "classifier_forward",
@@ -240,17 +232,14 @@ def forward(self, x):
         "set_optimizer",
     )
     def set_optimizer(self, optimizer_name, **kwargs):
-        self.optimizer_name = optimizer_name
-        self.optimizer_args = kwargs
+        super().set_optimizer(optimizer_name, **kwargs)
 
     @torchensemble_model_doc(
         """Set the attributes on scheduler for AdversarialTrainingClassifier.""",  # noqa: E501
         "set_scheduler",
     )
     def set_scheduler(self, scheduler_name, **kwargs):
-        self.scheduler_name = scheduler_name
-        self.scheduler_args = kwargs
-        self.use_scheduler_ = True
+        super().set_scheduler(scheduler_name, **kwargs)
 
     @_adversarial_training_model_doc(
         """Implementation on the training stage of AdversarialTrainingClassifier.""",  # noqa: E501
@@ -268,7 +257,7 @@ def fit(
     ):
 
         self._validate_parameters(epochs, epsilon, log_interval)
-        self.n_outputs = self._decide_n_outputs(train_loader, True)
+        self.n_outputs = self._decide_n_outputs(train_loader)
 
         # Instantiate a pool of base estimators, optimizers, and schedulers.
         estimators = []
@@ -404,10 +393,6 @@ def predict(self, X, return_numpy=True):
     "model",
 )
 class AdversarialTrainingRegressor(_BaseAdversarialTraining, BaseRegressor):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-        self.is_classification = False
-
     @torchensemble_model_doc(
         """Implementation on the data forwarding in AdversarialTrainingRegressor.""",  # noqa: E501
         "regressor_forward",
@@ -424,17 +409,14 @@ def forward(self, x):
         "set_optimizer",
     )
     def set_optimizer(self, optimizer_name, **kwargs):
-        self.optimizer_name = optimizer_name
-        self.optimizer_args = kwargs
+        super().set_optimizer(optimizer_name, **kwargs)
 
     @torchensemble_model_doc(
         """Set the attributes on scheduler for AdversarialTrainingRegressor.""",  # noqa: E501
         "set_scheduler",
     )
     def set_scheduler(self, scheduler_name, **kwargs):
-        self.scheduler_name = scheduler_name
-        self.scheduler_args = kwargs
-        self.use_scheduler_ = True
+        super().set_scheduler(scheduler_name, **kwargs)
 
     @_adversarial_training_model_doc(
         """Implementation on the training stage of AdversarialTrainingRegressor.""",  # noqa: E501
@@ -452,7 +434,7 @@ def fit(
     ):
 
         self._validate_parameters(epochs, epsilon, log_interval)
-        self.n_outputs = self._decide_n_outputs(train_loader, True)
+        self.n_outputs = self._decide_n_outputs(train_loader)
 
         # Instantiate a pool of base estimators, optimizers, and schedulers.
         estimators = []

torchensemble/bagging.py

@@ -115,17 +115,14 @@ def forward(self, x):
         "set_optimizer",
     )
     def set_optimizer(self, optimizer_name, **kwargs):
-        self.optimizer_name = optimizer_name
-        self.optimizer_args = kwargs
+        super().set_optimizer(optimizer_name, **kwargs)
 
     @torchensemble_model_doc(
         """Set the attributes on scheduler for BaggingClassifier.""",
         "set_scheduler",
     )
     def set_scheduler(self, scheduler_name, **kwargs):
-        self.scheduler_name = scheduler_name
-        self.scheduler_args = kwargs
-        self.use_scheduler_ = True
+        super().set_scheduler(scheduler_name, **kwargs)
 
     @torchensemble_model_doc(
         """Implementation on the training stage of BaggingClassifier.""", "fit"
@@ -141,7 +138,7 @@ def fit(
     ):
 
         self._validate_parameters(epochs, log_interval)
-        self.n_outputs = self._decide_n_outputs(train_loader, True)
+        self.n_outputs = self._decide_n_outputs(train_loader)
 
         # Instantiate a pool of base estimators, optimizers, and schedulers.
         estimators = []
@@ -289,17 +286,14 @@ def forward(self, x):
         "set_optimizer",
     )
     def set_optimizer(self, optimizer_name, **kwargs):
-        self.optimizer_name = optimizer_name
-        self.optimizer_args = kwargs
+        super().set_optimizer(optimizer_name, **kwargs)
 
     @torchensemble_model_doc(
         """Set the attributes on scheduler for BaggingRegressor.""",
         "set_scheduler",
     )
     def set_scheduler(self, scheduler_name, **kwargs):
-        self.scheduler_name = scheduler_name
-        self.scheduler_args = kwargs
-        self.use_scheduler_ = True
+        super().set_scheduler(scheduler_name, **kwargs)
 
     @torchensemble_model_doc(
         """Implementation on the training stage of BaggingRegressor.""", "fit"
@@ -315,7 +309,7 @@ def fit(
     ):
 
         self._validate_parameters(epochs, log_interval)
-        self.n_outputs = self._decide_n_outputs(train_loader, False)
+        self.n_outputs = self._decide_n_outputs(train_loader)
 
         # Instantiate a pool of base estimators, optimizers, and schedulers.
         estimators = []

torchensemble/fast_geometric.py

@@ -24,11 +24,7 @@
 from .utils.logging import get_tb_logger
 
 
-__all__ = [
-    "_BaseFastGeometric",
-    "FastGeometricClassifier",
-    "FastGeometricRegressor",
-]
+__all__ = ["FastGeometricClassifier", "FastGeometricRegressor"]
 
 
 __fit_doc = """
@@ -175,10 +171,6 @@ def set_scheduler(self, scheduler_name, **kwargs):
     """Implementation on the FastGeometricClassifier.""", "seq_model"
 )
 class FastGeometricClassifier(_BaseFastGeometric, BaseClassifier):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-        self.is_classification = True
-
     @torchensemble_model_doc(
         """Implementation on the data forwarding in FastGeometricClassifier.""",  # noqa: E501
         "classifier_forward",
@@ -227,9 +219,7 @@ def fit(
         save_dir=None,
     ):
         self._validate_parameters(epochs, log_interval)
-        self.n_outputs = self._decide_n_outputs(
-            train_loader, self.is_classification
-        )
+        self.n_outputs = self._decide_n_outputs(train_loader)
 
         # ====================================================================
         #                Train the dummy estimator (estimator_)
@@ -426,10 +416,6 @@ def predict(self, X, return_numpy=True):
     """Implementation on the FastGeometricRegressor.""", "seq_model"
 )
 class FastGeometricRegressor(_BaseFastGeometric, BaseRegressor):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-        self.is_classification = False
-
     @torchensemble_model_doc(
         """Implementation on the data forwarding in FastGeometricRegressor.""",  # noqa: E501
         "regressor_forward",
@@ -477,9 +463,7 @@ def fit(
         save_dir=None,
     ):
         self._validate_parameters(epochs, log_interval)
-        self.n_outputs = self._decide_n_outputs(
-            train_loader, self.is_classification
-        )
+        self.n_outputs = self._decide_n_outputs(train_loader)
 
         # ====================================================================
         #                Train the dummy estimator (estimator_)

torchensemble/fusion.py

@@ -49,17 +49,14 @@ def forward(self, x):
         "set_optimizer",
     )
     def set_optimizer(self, optimizer_name, **kwargs):
-        self.optimizer_name = optimizer_name
-        self.optimizer_args = kwargs
+        super().set_optimizer(optimizer_name, **kwargs)
 
     @torchensemble_model_doc(
         """Set the attributes on scheduler for FusionClassifier.""",
         "set_scheduler",
     )
     def set_scheduler(self, scheduler_name, **kwargs):
-        self.scheduler_name = scheduler_name
-        self.scheduler_args = kwargs
-        self.use_scheduler_ = True
+        super().set_scheduler(scheduler_name, **kwargs)
 
     @torchensemble_model_doc(
         """Implementation on the training stage of FusionClassifier.""", "fit"
@@ -78,7 +75,7 @@ def fit(
         for _ in range(self.n_estimators):
             self.estimators_.append(self._make_estimator())
         self._validate_parameters(epochs, log_interval)
-        self.n_outputs = self._decide_n_outputs(train_loader, True)
+        self.n_outputs = self._decide_n_outputs(train_loader)
         optimizer = set_module.set_optimizer(
             self, self.optimizer_name, **self.optimizer_args
         )
@@ -192,17 +189,14 @@ def forward(self, x):
         "set_optimizer",
     )
     def set_optimizer(self, optimizer_name, **kwargs):
-        self.optimizer_name = optimizer_name
-        self.optimizer_args = kwargs
+        super().set_optimizer(optimizer_name, **kwargs)
 
     @torchensemble_model_doc(
         """Set the attributes on scheduler for FusionRegressor.""",
         "set_scheduler",
     )
     def set_scheduler(self, scheduler_name, **kwargs):
-        self.scheduler_name = scheduler_name
-        self.scheduler_args = kwargs
-        self.use_scheduler_ = True
+        super().set_scheduler(scheduler_name, **kwargs)
 
     @torchensemble_model_doc(
         """Implementation on the training stage of FusionRegressor.""", "fit"
@@ -220,7 +214,7 @@ def fit(
         for _ in range(self.n_estimators):
             self.estimators_.append(self._make_estimator())
         self._validate_parameters(epochs, log_interval)
-        self.n_outputs = self._decide_n_outputs(train_loader, False)
+        self.n_outputs = self._decide_n_outputs(train_loader)
         optimizer = set_module.set_optimizer(
             self, self.optimizer_name, **self.optimizer_args
         )