diff --git a/pytorch_lightning/trainer/distrib_data_parallel.py b/pytorch_lightning/trainer/distrib_data_parallel.py index 75a60a45..abb2a32f 100644 --- a/pytorch_lightning/trainer/distrib_data_parallel.py +++ b/pytorch_lightning/trainer/distrib_data_parallel.py @@ -291,7 +291,7 @@ class TrainerDDPMixin(ABC): gpu_str = ','.join([str(x) for x in data_parallel_device_ids]) os.environ["CUDA_VISIBLE_DEVICES"] = gpu_str - log.info(f'VISIBLE GPUS: {os.environ["CUDA_VISIBLE_DEVICES"]}') + log.info(f'CUDA_VISIBLE_DEVICES: [{os.environ["CUDA_VISIBLE_DEVICES"]}]') def ddp_train(self, process_idx, model): """ diff --git a/pytorch_lightning/trainer/distrib_parts.py b/pytorch_lightning/trainer/distrib_parts.py index b0d7cae2..06905be9 100644 --- a/pytorch_lightning/trainer/distrib_parts.py +++ b/pytorch_lightning/trainer/distrib_parts.py @@ -628,7 +628,7 @@ def normalize_parse_gpu_string_input(s): if s == '-1': return -1 else: - return [int(x.strip()) for x in s.split(',')] + return [int(x.strip()) for x in s.split(',') if len(x) > 0] else: return s @@ -697,6 +697,10 @@ def parse_gpu_ids(gpus): then a misconfiguration exception is raised. """ + # nothing was passed into the GPUs argument + if callable(gpus): + return None + # Check that gpus param is None, Int, String or List check_gpus_data_type(gpus) diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py index ce8b9ca0..67cdb2d4 100644 --- a/pytorch_lightning/trainer/trainer.py +++ b/pytorch_lightning/trainer/trainer.py @@ -599,10 +599,25 @@ class Trainer( # TODO: get "help" from docstring :) for arg, arg_types, arg_default in (at for at in cls.get_init_arguments_and_types() if at[0] not in depr_arg_names): + for allowed_type in (at for at in allowed_types if at in arg_types): if allowed_type is bool: def allowed_type(x): return bool(distutils.util.strtobool(x)) + + if arg == 'gpus': + def allowed_type(x): + if ',' in x: + return str(x) + else: + return int(x) + + def arg_default(x): + if ',' in x: + return str(x) + else: + return int(x) + parser.add_argument( f'--{arg}', default=arg_default, diff --git a/tests/models/test_gpu.py b/tests/models/test_gpu.py index 38d519fb..69580fdd 100644 --- a/tests/models/test_gpu.py +++ b/tests/models/test_gpu.py @@ -259,6 +259,7 @@ def test_determine_root_gpu_device(gpus, expected_root_gpu): pytest.param('0', [0]), pytest.param('3', [3]), pytest.param('1, 3', [1, 3]), + pytest.param('2,', [2]), pytest.param('-1', list(range(PRETEND_N_OF_GPUS)), id="'-1' - use all gpus"), ]) def test_parse_gpu_ids(mocked_device_count, gpus, expected_gpu_ids): @@ -281,13 +282,6 @@ def test_parse_gpu_fail_on_unsupported_inputs(mocked_device_count, gpus): parse_gpu_ids(gpus) -@pytest.mark.gpus_param_tests -def test_parse_gpu_fail_on_empty_string(mocked_device_count): - # This currently results in a ValueError instead of MisconfigurationException - with pytest.raises(ValueError): - parse_gpu_ids('') - - @pytest.mark.gpus_param_tests @pytest.mark.parametrize("gpus", [[1, 2, 19], -1, '-1']) def test_parse_gpu_fail_on_non_existant_id(mocked_device_count_0, gpus):