Optimizer Frequencies logic, and new configure_optimizers

pytorch_lightning/core/lightning.py

@@ -914,10 +914,20 @@ def configure_optimizers(self) -> Union[
 
         If you don't define this method Lightning will automatically use Adam(lr=1e-3)
 
-        Return: any of these 3 options:
-            - Single optimizer
-            - List or Tuple - List of optimizers
-            - Two lists - The first list has multiple optimizers, the second a list of LR schedulers
+        Return: any of these 5 options:
+            - Single optimizer.
+            - List or Tuple - List of optimizers.
+            - Two lists - The first list has multiple optimizers, the second a list of LR schedulers.
+            - Dictionary, with an `optimizer` key and (optionally) a `lr_scheduler` key.
+            - Tuple of dictionaries as described, with an optional `frequency` key.
+
+        Note:
+            The `frequency` value is an int corresponding to the number of sequential batches
+            optimized with the specific optimizer. It should be given to none or to all of the optimizers.
+            There is difference between passing multiple optimizers in a list,
+            and passing multiple optimizers in dictionaries with a frequency of 1:
+            In the former case, all optimizers will operate on the given batch in each optimization step.
+            In the latter, only one optimizer will operate on the given batch at every step.
 
         Examples:
             .. code-block:: python
@@ -949,6 +959,18 @@ def configure_optimizers(self):
                     dis_sched = CosineAnnealing(discriminator_opt, T_max=10) # called every epoch
                     return [gen_opt, dis_opt], [gen_sched, dis_sched]
 
+                # example with optimizer frequencies
+                # see training procedure in `Improved Training of Wasserstein GANs`, Algorithm 1
+                # https://arxiv.org/abs/1704.00028
+                def configure_optimizers(self):
+                    gen_opt = Adam(self.model_gen.parameters(), lr=0.01)
+                    dis_opt = Adam(self.model_disc.parameters(), lr=0.02)
+                    n_critic = 5
+                    return (
+                        {'optimizer': dis_opt, 'frequency': n_critic},
+                        {'optimizer': gen_opt, 'frequency': 1}
+                    )
+
         Note:
 
             Some things to know:

pytorch_lightning/trainer/distrib_data_parallel.py

@@ -305,7 +305,8 @@ def ddp_train(self, gpu_idx, model):
 
         # CHOOSE OPTIMIZER
         # allow for lr schedulers as well
-        self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers())
+        self.optimizers, self.lr_schedulers, self.optimizer_frequencies = \
+            self.init_optimizers(model.configure_optimizers())
 
         # MODEL
         # copy model to each gpu

pytorch_lightning/trainer/distrib_parts.py

@@ -459,7 +459,8 @@ def single_gpu_train(self, model):
 
         # CHOOSE OPTIMIZER
         # allow for lr schedulers as well
-        self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers())
+        self.optimizers, self.lr_schedulers, self.optimizer_frequencies = \
+            self.init_optimizers(model.configure_optimizers())
 
         if self.use_amp:
             # An example
@@ -485,7 +486,8 @@ def tpu_train(self, tpu_core_idx, model):
 
         # CHOOSE OPTIMIZER
         # allow for lr schedulers as well
-        self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers())
+        self.optimizers, self.lr_schedulers, self.optimizer_frequencies = \
+            self.init_optimizers(model.configure_optimizers())
 
         # init 16 bit for TPU
         if self.precision == 16:
@@ -504,7 +506,8 @@ def dp_train(self, model):
 
         # CHOOSE OPTIMIZER
         # allow for lr schedulers as well
-        self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers())
+        self.optimizers, self.lr_schedulers, self.optimizer_frequencies = \
+            self.init_optimizers(model.configure_optimizers())
 
         model.cuda(self.root_gpu)
 

pytorch_lightning/trainer/trainer.py

@@ -3,7 +3,7 @@
 import sys
 import warnings
 from argparse import ArgumentParser
-from typing import Union, Optional, List, Dict, Tuple, Iterable, Any
+from typing import Union, Optional, List, Dict, Tuple, Iterable, Any, Sequence
 import distutils
 
 import torch
@@ -358,6 +358,7 @@ def __init__(
         self.disable_validation = False
         self.lr_schedulers = []
         self.optimizers = None
+        self.optimizer_frequencies = []
         self.global_step = 0
         self.current_epoch = 0
         self.total_batches = 0
@@ -714,7 +715,8 @@ def fit(
 
             # CHOOSE OPTIMIZER
             # allow for lr schedulers as well
-            self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers())
+            self.optimizers, self.lr_schedulers, self.optimizer_frequencies = \
+                self.init_optimizers(model.configure_optimizers())
 
             self.run_pretrain_routine(model)
 
@@ -760,31 +762,57 @@ def __attach_dataloaders(self, model, train_dataloader, val_dataloaders, test_da
 
     def init_optimizers(
             self,
-            optimizers: Union[Optimizer, Tuple[List, List], List[Optimizer], Tuple[Optimizer]]
-    ) -> Tuple[List, List]:
+            optim_conf: Union[Optimizer, Sequence[Optimizer], Dict, Sequence[Dict], Tuple[List, List]]
+    ) -> Tuple[List, List, List]:
 
         # single output, single optimizer
-        if isinstance(optimizers, Optimizer):
-            return [optimizers], []
+        if isinstance(optim_conf, Optimizer):
+            return [optim_conf], [], []
 
         # two lists, optimizer + lr schedulers
-        elif len(optimizers) == 2 and isinstance(optimizers[0], list):
-            optimizers, lr_schedulers = optimizers
+        elif isinstance(optim_conf, (list, tuple)) and len(optim_conf) == 2 and isinstance(optim_conf[0], list):
+            optimizers, lr_schedulers = optim_conf
             lr_schedulers = self.configure_schedulers(lr_schedulers)
-            return optimizers, lr_schedulers
+            return optimizers, lr_schedulers, []
+
+        # single dictionary
+        elif isinstance(optim_conf, dict):
+            optimizer = optim_conf["optimizer"]
+            lr_scheduler = optim_conf.get("lr_scheduler", [])
+            if lr_scheduler:
+                lr_schedulers = self.configure_schedulers([lr_scheduler])
+            return [optimizer], lr_schedulers, []
+
+        # multiple dictionaries
+        elif isinstance(optim_conf, (list, tuple)) and isinstance(optim_conf[0], dict):
+            optimizers = [opt_dict["optimizer"] for opt_dict in optim_conf]
+            # take only lr wif exists and ot they are defined - not None
+            lr_schedulers = [opt_dict["lr_scheduler"] for opt_dict in optim_conf if opt_dict.get("lr_scheduler")]
+            # take only freq wif exists and ot they are defined - not None
+            optimizer_frequencies = [opt_dict["frequency"] for opt_dict in optim_conf if opt_dict.get("frequency")]
+
+            # clean scheduler list
+            if lr_schedulers:
+                lr_schedulers = self.configure_schedulers(lr_schedulers)
+            # assert that if frequencies are present, they are given for all optimizers
+            if optimizer_frequencies and len(optimizer_frequencies) != len(optimizers):
+                raise ValueError("A frequency must be given to each optimizer.")
+            return optimizers, lr_schedulers, optimizer_frequencies
 
         # single list or tuple, multiple optimizer
-        elif isinstance(optimizers, (list, tuple)):
-            return optimizers, []
+        elif isinstance(optim_conf, (list, tuple)):
+            return list(optim_conf), [], []
 
         # unknown configuration
         else:
-            raise ValueError('Unknown configuration for model optimizers. Output'
-                             'from model.configure_optimizers() should either be:'
-                             '* single output, single torch.optim.Optimizer'
-                             '* single output, list of torch.optim.Optimizer'
-                             '* two outputs, first being a list of torch.optim.Optimizer',
-                             'second being a list of torch.optim.lr_scheduler')
+            raise ValueError(
+                'Unknown configuration for model optimizers. Output from `model.configure_optimizers()` should either be:'
+                ' * single output, single `torch.optim.Optimizer`'
+                ' * single output, list of `torch.optim.Optimizer`'
+                ' * single output, a dictionary with `optimizer` key (`torch.optim.Optimizer`)'
+                '    and an optional `lr_scheduler` key (`torch.optim.lr_scheduler`)'
+                ' * two outputs, first being a list of `torch.optim.Optimizer` second being a list of `torch.optim.lr_scheduler`'
+                ' * multiple outputs, dictionaries as described with an optional `frequency` key (int)')
 
     def configure_schedulers(self, schedulers: list):
         # Convert each scheduler into dict sturcture with relevant information
@@ -971,6 +999,7 @@ class _PatchDataLoader(object):
         dataloader: Dataloader object to return when called.
 
     """
+
     def __init__(self, dataloader: Union[List[DataLoader], DataLoader]):
         self.dataloader = dataloader
 

pytorch_lightning/trainer/training_loop.py

@@ -196,6 +196,7 @@ class TrainerTrainLoopMixin(ABC):
     total_batches: int
     truncated_bptt_steps: ...
     optimizers: ...
+    optimizer_frequencies: ...
     accumulate_grad_batches: int
     use_amp: bool
     track_grad_norm: ...
@@ -515,8 +516,7 @@ def run_training_batch(self, batch, batch_idx):
         for split_idx, split_batch in enumerate(splits):
             self.split_idx = split_idx
 
-            # call training_step once per optimizer
-            for opt_idx, optimizer in enumerate(self.optimizers):
+            for opt_idx, optimizer in self._get_optimizers_iterable():
                 # make sure only the gradients of the current optimizer's paramaters are calculated
                 # in the training step to prevent dangling gradients in multiple-optimizer setup.
                 if len(self.optimizers) > 1:
@@ -617,6 +617,19 @@ def optimizer_closure():
         self.callback_metrics.update({k: v for d in all_callback_metrics for k, v in d.items()})
 
         return 0, grad_norm_dic, all_log_metrics
+
+    def _get_optimizers_iterable(self):
+        if not self.optimizer_frequencies:
+            # call training_step once per optimizer
+            return list(enumerate(self.optimizers))
+
+        optimizer_freq_cumsum = np.cumsum(self.optimizer_frequencies)
+        optimizers_loop_length = optimizer_freq_cumsum[-1]
+        current_place_in_loop = self.total_batch_idx % optimizers_loop_length
+
+        # find optimzier index by looking for the first {item > current_place} in the cumsum list
+        opt_idx = np.argmax(optimizer_freq_cumsum > current_place_in_loop)
+        return [(opt_idx, self.optimizers[opt_idx])]
 
     def run_training_teardown(self):
         self.main_progress_bar.close()

tests/base/utils.py

@@ -82,7 +82,8 @@ def run_model_test(trainer_options, model, on_gpu=True):
     if trainer.use_ddp or trainer.use_ddp2:
         # on hpc this would work fine... but need to hack it for the purpose of the test
         trainer.model = pretrained_model
-        trainer.optimizers, trainer.lr_schedulers = trainer.init_optimizers(pretrained_model.configure_optimizers())
+        trainer.optimizers, trainer.lr_schedulers, trainer.optimizer_frequencies = \
+            trainer.init_optimizers(pretrained_model.configure_optimizers())
 
     # test HPC loading / saving
     trainer.hpc_save(save_dir, logger)

tests/models/test_gpu.py

@@ -96,30 +96,52 @@ def test_optimizer_return_options():
     # single optimizer
     opt_a = torch.optim.Adam(model.parameters(), lr=0.002)
     opt_b = torch.optim.SGD(model.parameters(), lr=0.002)
-    optim, lr_sched = trainer.init_optimizers(opt_a)
-    assert len(optim) == 1 and len(lr_sched) == 0
+    scheduler_a = torch.optim.lr_scheduler.StepLR(opt_a, 10)
+    scheduler_b = torch.optim.lr_scheduler.StepLR(opt_b, 10)
+
+    # single optimizer
+    optim, lr_sched, freq = trainer.init_optimizers(opt_a)
+    assert len(optim) == 1 and len(lr_sched) == 0 and len(freq) == 0
 
     # opt tuple
     opts = (opt_a, opt_b)
-    optim, lr_sched = trainer.init_optimizers(opts)
+    optim, lr_sched, freq = trainer.init_optimizers(opts)
     assert len(optim) == 2 and optim[0] == opts[0] and optim[1] == opts[1]
-    assert len(lr_sched) == 0
+    assert len(lr_sched) == 0 and len(freq) == 0
 
     # opt list
     opts = [opt_a, opt_b]
-    optim, lr_sched = trainer.init_optimizers(opts)
+    optim, lr_sched, freq = trainer.init_optimizers(opts)
     assert len(optim) == 2 and optim[0] == opts[0] and optim[1] == opts[1]
-    assert len(lr_sched) == 0
-
-    # opt tuple of lists
-    scheduler = torch.optim.lr_scheduler.StepLR(opt_a, 10)
-    opts = ([opt_a], [scheduler])
-    optim, lr_sched = trainer.init_optimizers(opts)
-    assert len(optim) == 1 and len(lr_sched) == 1
-    assert optim[0] == opts[0][0] and \
-        lr_sched[0] == dict(scheduler=scheduler, interval='epoch',
-                            frequency=1, reduce_on_plateau=False,
-                            monitor='val_loss')
+    assert len(lr_sched) == 0 and len(freq) == 0
+
+    # opt tuple of 2 lists
+    opts = ([opt_a], [scheduler_a])
+    optim, lr_sched, freq = trainer.init_optimizers(opts)
+    assert len(optim) == 1 and len(lr_sched) == 1 and len(freq) == 0
+    assert optim[0] == opts[0][0]
+    assert lr_sched[0] == dict(scheduler=scheduler_a, interval='epoch',
+                               frequency=1, reduce_on_plateau=False, monitor='val_loss')
+
+    # opt single dictionary
+    opts = {"optimizer": opt_a, "lr_scheduler": scheduler_a}
+    optim, lr_sched, freq = trainer.init_optimizers(opts)
+    assert len(optim) == 1 and len(lr_sched) == 1 and len(freq) == 0
+    assert optim[0] == opt_a
+    assert lr_sched[0] == dict(scheduler=scheduler_a, interval='epoch',
+                               frequency=1, reduce_on_plateau=False, monitor='val_loss')
+
+    # opt multiple dictionaries with frequencies
+    opts = (
+        {"optimizer": opt_a, "lr_scheduler": scheduler_a, "frequency": 1},
+        {"optimizer": opt_b, "lr_scheduler": scheduler_b, "frequency": 5},
+    )
+    optim, lr_sched, freq = trainer.init_optimizers(opts)
+    assert len(optim) == 2 and len(lr_sched) == 2 and len(freq) == 2
+    assert optim[0] == opt_a
+    assert lr_sched[0] == dict(scheduler=scheduler_a, interval='epoch',
+                               frequency=1, reduce_on_plateau=False, monitor='val_loss')
+    assert freq == [1, 5]
 
 
 def test_cpu_slurm_save_load(tmpdir):