[DOC] Improve docstrings for better readability

* [DOC] Improve docstrings

* [MNT] Update unit tests

* flake8 formatting

Author: xuyxu

torchensemble/_base.py

@@ -74,7 +74,7 @@ def __len__(self):
         """
         Return the number of base estimators in the ensemble. The real number
         of base estimators may not match `self.n_estimators` because of the
-        early stopping stage in several ensembles.
+        early stopping stage in several ensembles such as Gradient Boosting.
         """
         return len(self.estimators_)
 

torchensemble/_constants.py

@@ -83,7 +83,7 @@
 
         - If ``None``, the model will be saved in the current directory.
         - If not ``None``, the model will be saved in the specified
-          directory.
+          directory: ``save_dir``.
 """
 
 

torchensemble/adversarial_training.py

@@ -31,30 +31,32 @@
     Parameters
     ----------
     train_loader : torch.utils.data.DataLoader
-        A :mod:`DataLoader` container that contains the training data.
+        A :mod:`torch.utils.data.DataLoader` container that contains the
+        training data.
     epochs : int, default=100
         The number of training epochs.
-    epsilon : float, defaul=0.01
+    epsilon : float, default=0.01
         The step used to generate adversarial samples in the fast gradient
         sign method (FGSM), which should be in the range [0, 1].
     log_interval : int, default=100
-        The number of batches to wait before printting the training status.
+        The number of batches to wait before logging the training status.
     test_loader : torch.utils.data.DataLoader, default=None
-        A :mod:`DataLoader` container that contains the evaluating data.
+        A :mod:`torch.utils.data.DataLoader` container that contains the
+        evaluating data.
 
         - If ``None``, no validation is conducted after each training
           epoch.
         - If not ``None``, the ensemble will be evaluated on this
           dataloader after each training epoch.
     save_model : bool, default=True
-        Whether to save the model.
+        Specify whether to save the model parameters.
 
-        - If test_loader is ``None``, the ensemble containing
-          ``n_estimators`` base estimators will be saved.
+        - If test_loader is ``None``, the ensemble fully trained will be
+          saved.
         - If test_loader is not ``None``, the ensemble with the best
           validation performance will be saved.
     save_dir : string, default=None
-        Specify where to save the model.
+        Specify where to save the model parameters.
 
         - If ``None``, the model will be saved in the current directory.
         - If not ``None``, the model will be saved in the specified

torchensemble/gradient_boosting.py

@@ -31,9 +31,10 @@
     n_estimators : int
         The number of base estimators in the ensemble.
     estimator_args : dict, default=None
-        The dictionary of parameters used to instantiate base estimators.
+        The dictionary of hyper-parameters used to instantiate base
+        estimators (Optional).
     shrinkage_rate : float, default=1
-        The shrinkage rate in gradient boosting.
+        The shrinkage rate used in gradient boosting.
     cuda : bool, default=True
 
         - If ``True``, use GPU to train and evaluate the ensemble.
@@ -50,34 +51,36 @@
     Parameters
     ----------
     train_loader : torch.utils.data.DataLoader
-        A :mod:`DataLoader` container that contains the training data.
+        A :mod:`torch.utils.data.DataLoader` container that contains the
+        training data.
     epochs : int, default=100
-        The number of training epochs.
+        The number of training epochs per base estimator.
     log_interval : int, default=100
-        The number of batches to wait before printting the training status.
+        The number of batches to wait before logging the training status.
     test_loader : torch.utils.data.DataLoader, default=None
-        A :mod:`DataLoader` container that contains the evaluating data.
+        A :mod:`torch.utils.data.DataLoader` container that contains the
+        evaluating data.
 
-        - If ``None``, no validation is conducted after each training
-          epoch.
+        - If ``None``, no validation is conducted after each base
+          estimator being trained.
         - If not ``None``, the ensemble will be evaluated on this
-          dataloader after each training epoch.
+          dataloader after each base estimator being trained.
     early_stopping_rounds : int, default=2
         Specify the number of tolerant rounds for early stopping. When the
         validation performance of the ensemble does not improve after
         adding the base estimator fitted in current iteration, the internal
         counter on early stopping will increase by one. When the value of
         the internal counter reaches ``early_stopping_rounds``, the
-        training stage  will terminate early.
+        training stage  will terminate instantly.
     save_model : bool, default=True
-        Whether to save the model.
+        Specify whether to save the model parameters.
 
         - If test_loader is ``None``, the ensemble containing
           ``n_estimators`` base estimators will be saved.
         - If test_loader is not ``None``, the ensemble with the best
           validation performance will be saved.
     save_dir : string, default=None
-        Specify where to save the model.
+        Specify where to save the model parameters.
 
         - If ``None``, the model will be saved in the current directory.
         - If not ``None``, the model will be saved in the specified

torchensemble/snapshot_ensemble.py

@@ -44,23 +44,24 @@
     epochs : int, default=100
         The number of training epochs.
     log_interval : int, default=100
-        The number of batches to wait before printting the training status.
+        The number of batches to wait before logging the training status.
     test_loader : torch.utils.data.DataLoader, default=None
-        A :mod:`DataLoader` container that contains the evaluating data.
+        A :mod:`torch.utils.data.DataLoader` container that contains the
+        evaluating data.
 
-        - If ``None``, no validation is conducted after each training
-          epoch.
+        - If ``None``, no validation is conducted after each snapshot model
+          being generated.
         - If not ``None``, the ensemble will be evaluated on this
-          dataloader after each training epoch.
+          dataloader after each snapshot model being generated.
     save_model : bool, default=True
-        Whether to save the model.
+        Specify whether to save the model parameters.
 
-        - If test_loader is ``None``, the ensemble containing
+        - If test_loader is ``None``, the ensemble with
           ``n_estimators`` base estimators will be saved.
         - If test_loader is not ``None``, the ensemble with the best
           validation performance will be saved.
     save_dir : string, default=None
-        Specify where to save the model.
+        Specify where to save the model parameters.
 
         - If ``None``, the model will be saved in the current directory.
         - If not ``None``, the model will be saved in the specified

torchensemble/tests/test_set_optimizer.py

@@ -39,9 +39,12 @@ def test_set_optimizer_normal(optimizer_name):
 
 def test_set_optimizer_Unknown():
     model = MLP()
-    with pytest.raises(NotImplementedError) as excinfo:
+
+    err_msg = ("Unrecognized optimizer: {}, should be one of"
+               " {{Adadelta, Adagrad, Adam, AdamW, Adamax, ASGD, RMSprop,"
+               " Rprop, SGD}}.").format("Unknown")
+    with pytest.raises(NotImplementedError, match=err_msg):
         torchensemble.utils.set_module.set_optimizer(model, "Unknown")
-    assert "Unknown name of the optimizer" in str(excinfo.value)
 
 
 def test_update_lr():

torchensemble/tests/test_set_scheduler.py

@@ -79,6 +79,11 @@ def test_set_scheduler_ReduceLROnPlateau():
 def test_set_scheduler_Unknown():
     model = MLP()
     optimizer = torch.optim.Adam(model.parameters())
-    with pytest.raises(NotImplementedError) as excinfo:
+
+    err_msg = ("Unrecognized scheduler: {}, should be one of"
+               " {{LambdaLR, MultiplicativeLR, StepLR, MultiStepLR,"
+               " ExponentialLR, CosineAnnealingLR, ReduceLROnPlateau,"
+               " CyclicLR, OneCycleLR, CosineAnnealingWarmRestarts}}.")
+    err_msg = err_msg.format("Unknown")
+    with pytest.raises(NotImplementedError, match=err_msg):
         torchensemble.utils.set_module.set_scheduler(optimizer, "Unknown")
-    assert "Unknown name of the scheduler" in str(excinfo.value)

torchensemble/utils/set_module.py

@@ -28,7 +28,7 @@ def set_optimizer(model, optimizer_name, **kwargs):
     elif optimizer_name == "SGD":
         optimizer = torch.optim.SGD(model.parameters(), **kwargs)
     else:
-        msg = ("Unknown name of the optimizer {}, should be one of"
+        msg = ("Unrecognized optimizer: {}, should be one of"
                " {{Adadelta, Adagrad, Adam, AdamW, Adamax, ASGD, RMSprop,"
                " Rprop, SGD}}.")
         raise NotImplementedError(msg.format(optimizer_name))
@@ -77,7 +77,7 @@ def set_scheduler(optimizer, scheduler_name, **kwargs):
     elif scheduler_name == "ReduceLROnPlateau":
         scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, **kwargs)
     else:
-        msg = ("Unknown name of the scheduler {}, should be one of"
+        msg = ("Unrecognized scheduler: {}, should be one of"
                " {{LambdaLR, MultiplicativeLR, StepLR, MultiStepLR,"
                " ExponentialLR, CosineAnnealingLR, ReduceLROnPlateau,"
                " CyclicLR, OneCycleLR, CosineAnnealingWarmRestarts}}.")