[FIX] Fix the binding during parallelization (#33)

* [FIX] Fix the binnding problem on scheduler

* flake8 formatting

* Update CHANGELOG.rst

Author: xuyxu

CHANGELOG.rst

@@ -16,7 +16,7 @@ Changelog
 * |Enhancement| Improve the logging module | @zzzzwj
 * |API| Remove the input argument ``output_dim`` from all methods | @xuyxu
 * |Fix| Fix the bug in logging module when using multi-processing | @zzzzwj
-
+* |Fix| Fix the binding problem on scheduler and optimizer when using parallelization | @Alex-Medium and @xuyxi
 
 .. role:: raw-html(raw)
    :format: html

torchensemble/adversarial_training.py

@@ -12,6 +12,8 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+
+import warnings
 from joblib import Parallel, delayed
 
 from ._base import BaseModule, torchensemble_model_doc
@@ -81,6 +83,7 @@ def adddoc(cls):
 def _parallel_fit_per_epoch(train_loader,
                             epsilon,
                             estimator,
+                            cur_lr,
                             optimizer,
                             criterion,
                             idx,
@@ -95,6 +98,10 @@ def _parallel_fit_per_epoch(train_loader,
     out-of-memory error.
     """
 
+    if cur_lr:
+        # Parallelization corrupts the binding between optimizer and scheduler
+        set_module.update_lr(optimizer, cur_lr)
+
     for batch_idx, (data, target) in enumerate(train_loader):
 
         batch_size = data.size()[0]
@@ -246,11 +253,9 @@ def fit(self,
                                                        **self.optimizer_args))
 
         if self.use_scheduler_:
-            schedulers = []
-            for i in range(self.n_estimators):
-                schedulers.append(set_module.set_scheduler(optimizers[i],
-                                                           self.scheduler_name,
-                                                           **self.scheduler_args))  # noqa: E501
+            scheduler_ = set_module.set_scheduler(optimizers[0],
+                                                  self.scheduler_name,
+                                                  **self.scheduler_args)
 
         # Utils
         criterion = nn.CrossEntropyLoss()
@@ -271,6 +276,11 @@ def _forward(estimators, data):
             for epoch in range(epochs):
                 self.train()
 
+                if self.use_scheduler_:
+                    cur_lr = scheduler_.get_last_lr()[0]
+                else:
+                    cur_lr = None
+
                 if self.n_jobs and self.n_jobs > 1:
                     msg = "Parallelization on the training epoch: {:03d}"
                     self.logger.info(msg.format(epoch))
@@ -279,6 +289,7 @@ def _forward(estimators, data):
                         train_loader,
                         epsilon,
                         estimator,
+                        cur_lr,
                         optimizer,
                         criterion,
                         idx,
@@ -323,9 +334,14 @@ def _forward(estimators, data):
                         self.logger.info(msg.format(epoch, acc, best_acc))
 
                 # Update the scheduler
-                if self.use_scheduler_:
-                    for i in range(self.n_estimators):
-                        schedulers[i].step()
+                with warnings.catch_warnings():
+
+                    # UserWarning raised by PyTorch is ignored because
+                    # scheduler does not have a real effect on the optimier.
+                    warnings.simplefilter("ignore", UserWarning)
+
+                    if self.use_scheduler_:
+                        scheduler_.step()
 
         self.estimators_ = nn.ModuleList()
         self.estimators_.extend(estimators)
@@ -413,11 +429,9 @@ def fit(self,
                                                        **self.optimizer_args))
 
         if self.use_scheduler_:
-            schedulers = []
-            for i in range(self.n_estimators):
-                schedulers.append(set_module.set_scheduler(optimizers[i],
-                                                           self.scheduler_name,
-                                                           **self.scheduler_args))  # noqa: E501
+            scheduler_ = set_module.set_scheduler(optimizers[0],
+                                                  self.scheduler_name,
+                                                  **self.scheduler_args)
 
         # Utils
         criterion = nn.MSELoss()
@@ -437,6 +451,11 @@ def _forward(estimators, data):
             for epoch in range(epochs):
                 self.train()
 
+                if self.use_scheduler_:
+                    cur_lr = scheduler_.get_last_lr()[0]
+                else:
+                    cur_lr = None
+
                 if self.n_jobs and self.n_jobs > 1:
                     msg = "Parallelization on the training epoch: {:03d}"
                     self.logger.info(msg.format(epoch))
@@ -445,6 +464,7 @@ def _forward(estimators, data):
                         train_loader,
                         epsilon,
                         estimator,
+                        cur_lr,
                         optimizer,
                         criterion,
                         idx,
@@ -486,9 +506,11 @@ def _forward(estimators, data):
                         self.logger.info(msg.format(epoch, mse, best_mse))
 
                 # Update the scheduler
-                if self.use_scheduler_:
-                    for i in range(self.n_estimators):
-                        schedulers[i].step()
+                with warnings.catch_warnings():
+                    warnings.simplefilter("ignore", UserWarning)
+
+                    if self.use_scheduler_:
+                        scheduler_.step()
 
         self.estimators_ = nn.ModuleList()
         self.estimators_.extend(estimators)

torchensemble/bagging.py

@@ -10,6 +10,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
+import warnings
 from joblib import Parallel, delayed
 
 from ._base import BaseModule, torchensemble_model_doc
@@ -24,6 +25,7 @@
 
 def _parallel_fit_per_epoch(train_loader,
                             estimator,
+                            cur_lr,
                             optimizer,
                             criterion,
                             idx,
@@ -38,6 +40,10 @@ def _parallel_fit_per_epoch(train_loader,
     out-of-memory error.
     """
 
+    if cur_lr:
+        # Parallelization corrupts the binding between optimizer and scheduler
+        set_module.update_lr(optimizer, cur_lr)
+
     for batch_idx, (data, target) in enumerate(train_loader):
 
         batch_size = data.size(0)
@@ -134,11 +140,9 @@ def fit(self,
                                                        **self.optimizer_args))
 
         if self.use_scheduler_:
-            schedulers = []
-            for i in range(self.n_estimators):
-                schedulers.append(set_module.set_scheduler(optimizers[i],
-                                                           self.scheduler_name,
-                                                           **self.scheduler_args))  # noqa: E501
+            scheduler_ = set_module.set_scheduler(optimizers[0],
+                                                  self.scheduler_name,
+                                                  **self.scheduler_args)
 
         # Utils
         criterion = nn.CrossEntropyLoss()
@@ -159,13 +163,19 @@ def _forward(estimators, data):
             for epoch in range(epochs):
                 self.train()
 
+                if self.use_scheduler_:
+                    cur_lr = scheduler_.get_last_lr()[0]
+                else:
+                    cur_lr = None
+
                 if self.n_jobs and self.n_jobs > 1:
                     msg = "Parallelization on the training epoch: {:03d}"
                     self.logger.info(msg.format(epoch))
 
                 rets = parallel(delayed(_parallel_fit_per_epoch)(
                         train_loader,
                         estimator,
+                        cur_lr,
                         optimizer,
                         criterion,
                         idx,
@@ -210,9 +220,14 @@ def _forward(estimators, data):
                         self.logger.info(msg.format(epoch, acc, best_acc))
 
                 # Update the scheduler
-                if self.use_scheduler_:
-                    for i in range(self.n_estimators):
-                        schedulers[i].step()
+                with warnings.catch_warnings():
+
+                    # UserWarning raised by PyTorch is ignored because
+                    # scheduler does not have a real effect on the optimier.
+                    warnings.simplefilter("ignore", UserWarning)
+
+                    if self.use_scheduler_:
+                        scheduler_.step()
 
         self.estimators_ = nn.ModuleList()
         self.estimators_.extend(estimators)
@@ -294,11 +309,9 @@ def fit(self,
                                                        **self.optimizer_args))
 
         if self.use_scheduler_:
-            schedulers = []
-            for i in range(self.n_estimators):
-                schedulers.append(set_module.set_scheduler(optimizers[i],
-                                                           self.scheduler_name,
-                                                           **self.scheduler_args))  # noqa: E501
+            scheduler_ = set_module.set_scheduler(optimizers[0],
+                                                  self.scheduler_name,
+                                                  **self.scheduler_args)
 
         # Utils
         criterion = nn.MSELoss()
@@ -318,13 +331,19 @@ def _forward(estimators, data):
             for epoch in range(epochs):
                 self.train()
 
+                if self.use_scheduler_:
+                    cur_lr = scheduler_.get_last_lr()[0]
+                else:
+                    cur_lr = None
+
                 if self.n_jobs and self.n_jobs > 1:
                     msg = "Parallelization on the training epoch: {:03d}"
                     self.logger.info(msg.format(epoch))
 
                 rets = parallel(delayed(_parallel_fit_per_epoch)(
                         train_loader,
                         estimator,
+                        cur_lr,
                         optimizer,
                         criterion,
                         idx,
@@ -366,9 +385,11 @@ def _forward(estimators, data):
                         self.logger.info(msg.format(epoch, mse, best_mse))
 
                 # Update the scheduler
-                if self.use_scheduler_:
-                    for i in range(self.n_estimators):
-                        schedulers[i].step()
+                with warnings.catch_warnings():
+                    warnings.simplefilter("ignore", UserWarning)
+
+                    if self.use_scheduler_:
+                        scheduler_.step()
 
         self.estimators_ = nn.ModuleList()
         self.estimators_.extend(estimators)

torchensemble/tests/test_set_optimizer.py

@@ -42,3 +42,29 @@ def test_set_optimizer_Unknown():
     with pytest.raises(NotImplementedError) as excinfo:
         torchensemble.utils.set_module.set_optimizer(model, "Unknown")
     assert "Unknown name of the optimizer" in str(excinfo.value)
+
+
+def test_update_lr():
+    cur_lr = 1e-4
+    model = MLP()
+    optimizer = torchensemble.utils.set_module.set_optimizer(model,
+                                                             "Adam",
+                                                             lr=1e-3)
+
+    optimizer = torchensemble.utils.set_module.update_lr(optimizer, cur_lr)
+
+    for group in optimizer.param_groups:
+        assert group["lr"] == cur_lr
+
+
+def test_update_lr_invalid():
+    cur_lr = 0
+    model = MLP()
+    optimizer = torchensemble.utils.set_module.set_optimizer(model,
+                                                             "Adam",
+                                                             lr=1e-3)
+
+    err_msg = ("The learning rate should be strictly positive, but got"
+               " {} instead.").format(cur_lr)
+    with pytest.raises(ValueError, match=err_msg):
+        torchensemble.utils.set_module.update_lr(optimizer, cur_lr)

torchensemble/utils/set_module.py

@@ -36,6 +36,24 @@ def set_optimizer(model, optimizer_name, **kwargs):
     return optimizer
 
 
+def update_lr(optimizer, lr):
+    """
+    Manually update the learning rate of the optimizer. This function is used
+    when the parallelization corrupts the bindings between the optimizer and
+    the scheduler.
+    """
+
+    if not lr > 0:
+        msg = ("The learning rate should be strictly positive, but got"
+               " {} instead.")
+        raise ValueError(msg.format(lr))
+
+    for group in optimizer.param_groups:
+        group["lr"] = lr
+
+    return optimizer
+
+
 def set_scheduler(optimizer, scheduler_name, **kwargs):
     """
     Set the scheduler on learning rate for the optimizer.