diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py index 6d4e4e3f..7d20b5c2 100644 --- a/pytorch_lightning/models/trainer.py +++ b/pytorch_lightning/models/trainer.py @@ -12,6 +12,7 @@ import torch from torch.utils.data.distributed import DistributedSampler import torch.multiprocessing as mp import torch.distributed as dist +from torch.optim.optimizer import Optimizer from pytorch_lightning.root_module.root_module import LightningModule from pytorch_lightning.root_module.memory import get_gpu_memory_map @@ -532,12 +533,15 @@ If you want each process to load the full dataset, ignore this warning. task = int(os.environ['SLURM_LOCALID']) self.ddp_train(task, model) else: - msg = """ -You requested %(nb_gpus)s GPUs but launched %(nb_tasks)s slurm tasks. -We will launch %(nb_gpus)s processes for you. -We recommend you let slurm manage the processes by setting: --ntasks-per-node=%(nb_gpus)s -If you're not using SLURM, ignore this message! -""" % {'nb_gpus': self.nb_requested_gpus, 'nb_tasks': self.nb_slurm_tasks} + nb_gpus = self.nb_requested_gpus + nb_tasks = self.nb_slurm_tasks + msg = f""" + You requested {nb_gpus}s GPUs but launched {nb_tasks}s slurm tasks. + We will launch {nb_gpus}s processes for you. + We recommend you let slurm manage the processes by setting: + --ntasks-per-node={nb_gpus}s + If you're not using SLURM, ignore this message! + """ warnings.warn(msg) mp.spawn(self.ddp_train, nprocs=len(self.data_parallel_device_ids), args=(model, )) @@ -558,9 +562,7 @@ If you're not using SLURM, ignore this message! # CHOOSE OPTIMIZER # allow for lr schedulers as well - self.optimizers = model.configure_optimizers() - if len(self.optimizers) == 2 and type(self.optimizers[0]) is list: - self.optimizers, self.lr_schedulers = self.optimizers + self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers()) self.__run_pretrain_routine(model) @@ -568,12 +570,25 @@ If you're not using SLURM, ignore this message! # used for testing or when we need to know that training succeeded return 1 + def init_optimizers(self, optimizers): + + # single optimizer + if isinstance(optimizers, Optimizer): + return [optimizers], [] + + # two lists + elif len(optimizers) == 2 and isinstance(optimizers[0], list): + optimizers, lr_schedulers = optimizers + return optimizers, lr_schedulers + + # single list or tuple + elif isinstance(optimizers, list) or isinstance(optimizers, tuple): + return optimizers, [] + def __single_gpu_train(self, model): # CHOOSE OPTIMIZER # allow for lr schedulers as well - self.optimizers = model.configure_optimizers() - if len(self.optimizers) == 2: - self.optimizers, self.lr_schedulers = self.optimizers + self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers()) model.cuda(self.data_parallel_device_ids[0]) @@ -590,20 +605,18 @@ If you're not using SLURM, ignore this message! # CHOOSE OPTIMIZER # allow for lr schedulers as well - self.optimizers = model.configure_optimizers() - if len(self.optimizers) == 2: - self.optimizers, self.lr_schedulers = self.optimizers + self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers()) model.cuda(self.data_parallel_device_ids[0]) # check for this bug (amp + dp + !01 doesn't work) # https://github.com/NVIDIA/apex/issues/227 if self.use_dp and self.use_amp: - m = """ -Amp level %r with DataParallel is not supported. -See this note from NVIDIA for more info: https://github.com/NVIDIA/apex/issues/227. -We recommend you switch to ddp if you want to use amp -""" % self.amp_level + m = f""" + Amp level {self.amp_level} with DataParallel is not supported. + See this note from NVIDIA for more info: https://github.com/NVIDIA/apex/issues/227. + We recommend you switch to ddp if you want to use amp + """ raise MisconfigurationException(m) model = LightningDataParallel(model, device_ids=self.data_parallel_device_ids) @@ -650,9 +663,7 @@ We recommend you switch to ddp if you want to use amp # CHOOSE OPTIMIZER # allow for lr schedulers as well - self.optimizers = model.configure_optimizers() - if len(self.optimizers) == 2: - self.optimizers, self.lr_schedulers = self.optimizers + self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers()) # MODEL # copy model to each gpu diff --git a/pytorch_lightning/testing/lm_test_module.py b/pytorch_lightning/testing/lm_test_module.py index ce0e7c3a..e56b6179 100644 --- a/pytorch_lightning/testing/lm_test_module.py +++ b/pytorch_lightning/testing/lm_test_module.py @@ -195,7 +195,7 @@ class LightningTestModel(LightningModule): optimizer = optim.Adam(self.parameters(), lr=self.hparams.learning_rate) # test returning only 1 list instead of 2 - return [optimizer] + return optimizer def __dataloader(self, train): # init data generators diff --git a/tests/test_models.py b/tests/test_models.py index b7722584..4813c542 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -26,6 +26,37 @@ np.random.seed(SEED) # ------------------------------------------------------------------------ # TESTS # ------------------------------------------------------------------------ + +def test_optimizer_return_options(): + + trainer = Trainer() + model, hparams = get_model() + + # 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 + + # opt tuple + opts = (opt_a, opt_b) + optim, lr_sched = trainer.init_optimizers(opts) + assert len(optim) == 2 and optim[0] == opts[0] and optim[1] == opts[1] + assert len(lr_sched) == 0 + + # opt list + opts = [opt_a, opt_b] + optim, lr_sched = 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 + opts = ([opt_a], ['lr_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] == 'lr_scheduler' + + def test_single_gpu_batch_parse(): if not torch.cuda.is_available(): warnings.warn('test_amp_gpu_ddp cannot run.'