Fix Horovod distributed backend to set the root_gpu property (#1669)

* params

* drop acc

* Fix Horovod distributed backend to set the root_gpu

* Fixed test

* Fixed tests

* Fixed lint

* Set root_gpu during initialization

* chlog

Co-authored-by: Jirka <jirka.borovec@seznam.cz>
This commit is contained in:
Travis Addair
2020-05-01 14:13:35 -04:00
committed by GitHub
co-authored by Jirka
parent 34bc149359
commit 2950f66983
5 changed files with 30 additions and 13 deletions
@@ -27,12 +27,14 @@ PATH_HERE = os.path.abspath(os.path.dirname(__file__))
PATH_ROOT = os.path.join(PATH_HERE, '..', '..', '..', '..')
sys.path.insert(0, os.path.abspath(PATH_ROOT))
from pytorch_lightning import Trainer # noqa: E402
from pytorch_lightning.callbacks import ModelCheckpoint # noqa: E402
import tests.base.utils as tutils # noqa: E402
parser = argparse.ArgumentParser()
parser.add_argument('--trainer-options', required=True)
parser.add_argument('--on-gpu', action='store_true', default=False)
def run_test_from_config(trainer_options):
@@ -44,11 +46,15 @@ def run_test_from_config(trainer_options):
trainer_options['checkpoint_callback'] = ModelCheckpoint(ckpt_path)
model, hparams = tutils.get_default_model()
tutils.run_model_test(trainer_options, model, version=0, with_hpc=False)
tutils.run_model_test(trainer_options, model, on_gpu=args.on_gpu, version=0, with_hpc=False)
# Horovod should be initialized following training. If not, this will raise an exception.
assert hvd.size() == 2
if args.on_gpu:
# Test the root_gpu property
assert Trainer(gpus=1, distributed_backend='horovod', max_epochs=1).root_gpu == hvd.local_rank()
if __name__ == "__main__":
args = parser.parse_args()