diff --git a/.run_local_tests.sh b/.run_local_tests.sh new file mode 100644 index 00000000..48cb5acf --- /dev/null +++ b/.run_local_tests.sh @@ -0,0 +1,2 @@ +rm -rf tests/save_dir* +coverage run --source pytorch_lightning -m py.test pytorch_lightning tests examples -v --doctest-modules \ No newline at end of file diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py index 375d3328..454ee6f0 100644 --- a/pytorch_lightning/trainer/trainer.py +++ b/pytorch_lightning/trainer/trainer.py @@ -302,6 +302,8 @@ class Trainer(TrainerIO): return # single GPU case + # in single gpu case we allow ddp so we can train on multiple + # nodes, 1 gpu per node if self.num_gpus == 1: self.single_gpu = True @@ -354,6 +356,14 @@ class Trainer(TrainerIO): # likely not on slurm, so set the slurm managed flag to false self.is_slurm_managing_tasks = False + # used for tests only, set this flag to simulate slurm managing a task + try: + should_fake = int(os.environ['FAKE_SLURM_MANAGING_TASKS']) + if should_fake: + self.is_slurm_managing_tasks = True + except Exception as e: + pass + def __set_nvidia_flags(self, is_slurm_managing_tasks, data_parallel_device_ids): if data_parallel_device_ids is None: return @@ -929,7 +939,7 @@ class Trainer(TrainerIO): # run all epochs for epoch_nb in range(self.current_epoch, self.max_nb_epochs): # set seed for distributed sampler (enables shuffling for each epoch) - if self.use_ddp: + if self.use_ddp and hasattr(self.train_dataloader.sampler, 'set_epoch'): self.train_dataloader.sampler.set_epoch(epoch_nb) # get model diff --git a/tests/test_models.py b/tests/test_models.py index 3c82492a..62963e2c 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -7,6 +7,7 @@ import pytest import numpy as np import torch from test_tube import Experiment +import time # sys.path += [os.path.abspath('..'), os.path.abspath('../..')] from pytorch_lightning import Trainer @@ -39,6 +40,53 @@ np.random.seed(SEED) # ------------------------------------------------------------------------ # TESTS # ------------------------------------------------------------------------ +def test_running_test_pretrained_model_ddp(): + """Verify test() on pretrained model""" + if not can_run_gpu_test(): + return + + hparams = get_hparams() + model = LightningTestModel(hparams) + + save_dir = init_save_dir() + + # exp file to get meta + exp = get_exp(False) + exp.argparse(hparams) + exp.save() + + # exp file to get weights + checkpoint = ModelCheckpoint(save_dir) + + trainer_options = dict( + show_progress_bar=False, + max_nb_epochs=1, + train_percent_check=0.4, + val_percent_check=0.2, + checkpoint_callback=checkpoint, + experiment=exp, + gpus=2, + distributed_backend='ddp' + ) + + # fit model + trainer = Trainer(**trainer_options) + result = trainer.fit(model) + + # correct result and ok accuracy + assert result == 1, 'training failed to complete' + pretrained_model = load_model(exp, save_dir, on_gpu=True, module_class=LightningTestModel) + + # run test set + new_trainer = Trainer(**trainer_options) + new_trainer.test(pretrained_model) + + run_prediction(model.test_dataloader, pretrained_model) + + # test we have good test accuracy + clear_save_dir() + + def test_dp_resume(): """ Make sure DP continues training correctly @@ -122,53 +170,6 @@ def test_dp_resume(): clear_save_dir() -def test_running_test_pretrained_model_ddp(): - """Verify test() on pretrained model""" - if not can_run_gpu_test(): - return - - hparams = get_hparams() - model = LightningTestModel(hparams) - - save_dir = init_save_dir() - - # exp file to get meta - exp = get_exp(False) - exp.argparse(hparams) - exp.save() - - # exp file to get weights - checkpoint = ModelCheckpoint(save_dir) - - trainer_options = dict( - show_progress_bar=False, - max_nb_epochs=1, - train_percent_check=0.4, - val_percent_check=0.2, - checkpoint_callback=checkpoint, - experiment=exp, - gpus=[0, 1], - distributed_backend='ddp' - ) - - # fit model - trainer = Trainer(**trainer_options) - result = trainer.fit(model) - - # correct result and ok accuracy - assert result == 1, 'training failed to complete' - pretrained_model = load_model(exp, save_dir, on_gpu=True, module_class=LightningTestModel) - - # run test set - new_trainer = Trainer(**trainer_options) - new_trainer.test(pretrained_model) - - run_prediction(model.test_dataloader, pretrained_model) - - # test we have good test accuracy - clear_save_dir() - - def test_running_test_after_fitting(): """Verify test() on fitted model""" hparams = get_hparams() @@ -249,6 +250,28 @@ def test_running_test_without_val(): clear_save_dir() +def test_multi_gpu_model_ddp(): + """ + Make sure DDP works + :return: + """ + if not can_run_gpu_test(): + return + + os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0]) + model, hparams = get_model() + trainer_options = dict( + show_progress_bar=False, + max_nb_epochs=1, + train_percent_check=0.4, + val_percent_check=0.2, + gpus=2, + distributed_backend='ddp' + ) + + run_gpu_model_test(trainer_options, model, hparams) + + def test_running_test_pretrained_model(): """Verify test() on pretrained model""" hparams = get_hparams() @@ -316,7 +339,7 @@ def test_running_test_pretrained_model_dp(): val_percent_check=0.2, checkpoint_callback=checkpoint, experiment=exp, - gpus=[0, 1], + gpus=2, distributed_backend='dp' ) @@ -408,28 +431,6 @@ def test_gradient_accumulation_scheduling(): trainer.fit(model) -def test_multi_gpu_model_ddp(): - """ - Make sure DDP works - :return: - """ - if not can_run_gpu_test(): - return - - os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0]) - model, hparams = get_model() - trainer_options = dict( - show_progress_bar=False, - max_nb_epochs=1, - train_percent_check=0.4, - val_percent_check=0.2, - gpus=[0, 1], - distributed_backend='ddp' - ) - - run_gpu_model_test(trainer_options, model, hparams) - - def test_optimizer_return_options(): trainer = Trainer() @@ -669,7 +670,7 @@ def test_amp_single_gpu(): show_progress_bar=True, max_nb_epochs=1, gpus=1, - distributed_backend='ddp', + distributed_backend='dp', use_amp=True ) @@ -978,79 +979,6 @@ def test_model_freeze_unfreeze(): model.unfreeze() -def test_amp_gpu_ddp_slurm_managed(): - """ - Make sure DDP + AMP work - :return: - """ - if not can_run_gpu_test(): - return - - # simulate setting slurm flags - os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0]) - os.environ['SLURM_LOCALID'] = str(0) - - hparams = get_hparams() - model = LightningTestModel(hparams) - - trainer_options = dict( - show_progress_bar=True, - max_nb_epochs=1, - gpus=[0], - distributed_backend='ddp', - use_amp=True - ) - - save_dir = init_save_dir() - - # exp file to get meta - exp = get_exp(False) - exp.argparse(hparams) - exp.save() - - # exp file to get weights - checkpoint = ModelCheckpoint(save_dir) - - # add these to the trainer options - trainer_options['checkpoint_callback'] = checkpoint - trainer_options['experiment'] = exp - - # fit model - trainer = Trainer(**trainer_options) - trainer.is_slurm_managing_tasks = True - result = trainer.fit(model) - - # correct result and ok accuracy - assert result == 1, 'amp + ddp model failed to complete' - - # test root model address - assert trainer.resolve_root_node_address('abc') == 'abc' - assert trainer.resolve_root_node_address('abc[23]') == 'abc23' - assert trainer.resolve_root_node_address('abc[23-24]') == 'abc23' - assert trainer.resolve_root_node_address('abc[23-24, 45-40, 40]') == 'abc23' - - # test model loading with a map_location - pretrained_model = load_model(exp, save_dir, True) - - # test model preds - run_prediction(model.test_dataloader, pretrained_model) - - if trainer.use_ddp: - # 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 = pretrained_model.configure_optimizers() - - # test HPC loading / saving - trainer.hpc_save(save_dir, exp) - trainer.hpc_load(save_dir, on_gpu=True) - - # test freeze on gpu - model.freeze() - model.unfreeze() - - clear_save_dir() - - def test_cpu_model_with_amp(): """ Make sure model trains on CPU @@ -1221,7 +1149,7 @@ def test_ddp_sampler_error(): experiment=exp, show_progress_bar=False, max_nb_epochs=1, - gpus=[0, 1], + gpus=2, distributed_backend='ddp', use_amp=True ) @@ -1300,6 +1228,79 @@ def test_multiple_test_dataloader(): trainer.test() +def test_amp_gpu_ddp_slurm_managed(): + """ + Make sure DDP + AMP work + :return: + """ + if not can_run_gpu_test(): + return + + # simulate setting slurm flags + os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0]) + os.environ['SLURM_LOCALID'] = str(0) + os.environ['FAKE_SLURM_MANAGING_TASKS'] = str(1) + + hparams = get_hparams() + model = LightningTestModel(hparams) + + trainer_options = dict( + show_progress_bar=True, + max_nb_epochs=1, + gpus=2, + distributed_backend='ddp', + use_amp=True + ) + + save_dir = init_save_dir() + + # exp file to get meta + exp = get_exp(False) + exp.argparse(hparams) + exp.save() + + # exp file to get weights + checkpoint = ModelCheckpoint(save_dir) + + # add these to the trainer options + trainer_options['checkpoint_callback'] = checkpoint + trainer_options['experiment'] = exp + + # fit model + trainer = Trainer(**trainer_options) + result = trainer.fit(model) + + # correct result and ok accuracy + assert result == 1, 'amp + ddp model failed to complete' + + # test root model address + assert trainer.resolve_root_node_address('abc') == 'abc' + assert trainer.resolve_root_node_address('abc[23]') == 'abc23' + assert trainer.resolve_root_node_address('abc[23-24]') == 'abc23' + assert trainer.resolve_root_node_address('abc[23-24, 45-40, 40]') == 'abc23' + + # test model loading with a map_location + pretrained_model = load_model(exp, save_dir, True) + + # test model preds + run_prediction(model.test_dataloader, pretrained_model) + + if trainer.use_ddp: + # 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 = pretrained_model.configure_optimizers() + + # test HPC loading / saving + trainer.hpc_save(save_dir, exp) + trainer.hpc_load(save_dir, on_gpu=True) + + # test freeze on gpu + model.freeze() + model.unfreeze() + + clear_save_dir() + + # ------------------------------------------------------------------------ # UTILS # ------------------------------------------------------------------------ @@ -1462,13 +1463,14 @@ def assert_ok_test_acc(trainer): def can_run_gpu_test(): if not torch.cuda.is_available(): - warnings.warn('test_multi_gpu_model_ddp cannot run.' + warnings.warn('GPU test cannot run.' ' Rerun on a GPU node to run this test') return False if not torch.cuda.device_count() > 1: - warnings.warn('test_multi_gpu_model_ddp cannot run.' + warnings.warn('GPU test cannot run.' ' Rerun on a node with 2+ GPUs to run this test') return False + return True