From b973e7ce543e612a804aec5fe65a108728485261 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 26 Nov 2020 20:54:43 +0100 Subject: [PATCH] Enable use of amo in pipeline and CLI --- src/pykeen/models/cli/builders.py | 3 +++ src/pykeen/models/cli/options.py | 5 +++++ src/pykeen/pipeline.py | 14 +++++++++++--- tests/test_models.py | 10 ++++++++++ 4 files changed, 29 insertions(+), 3 deletions(-) diff --git a/src/pykeen/models/cli/builders.py b/src/pykeen/models/cli/builders.py index fd266b449a..cbb7e90b0c 100644 --- a/src/pykeen/models/cli/builders.py +++ b/src/pykeen/models/cli/builders.py @@ -101,6 +101,7 @@ def _decorate_model_kwargs(command: click.Command) -> click.Command: @options.optimizer_option @regularizer_option @options.training_loop_option + @options.automatic_memory_optimization_option @options.number_epochs_option @options.batch_size_option @options.learning_rate_option @@ -128,6 +129,7 @@ def main( mlflow_tracking_uri, title, dataset, + automatic_memory_optimization, training_triples_factory, testing_triples_factory, validation_triples_factory, @@ -180,6 +182,7 @@ def main( title=title, ), random_seed=random_seed, + automatic_memory_optimization=automatic_memory_optimization, ) if not silent: diff --git a/src/pykeen/models/cli/options.py b/src/pykeen/models/cli/options.py index 788fb0d8f9..32a0437516 100644 --- a/src/pykeen/models/cli/options.py +++ b/src/pykeen/models/cli/options.py @@ -137,6 +137,11 @@ def triples_factory_callback(_, __, path: Optional[str]) -> Optional[TriplesFact default=_get_default(get_training_loop_cls, suffix=_TRAINING_LOOP_SUFFIX), show_default=True, ) +automatic_memory_optimization_option = click.option( + '--automatic-memory-optimization/--no-automatic-memory-optimization', + default=True, + show_default=True, +) stopper_option = click.option( '--stopper', type=click.Choice(list(stoppers)), diff --git a/src/pykeen/pipeline.py b/src/pykeen/pipeline.py index 472c9c76c8..e429548488 100644 --- a/src/pykeen/pipeline.py +++ b/src/pykeen/pipeline.py @@ -738,6 +738,7 @@ def pipeline( # noqa: C901 result_tracker: Union[None, str, Type[ResultTracker]] = None, result_tracker_kwargs: Optional[Mapping[str, Any]] = None, # Misc + automatic_memory_optimization: bool = True, metadata: Optional[Dict[str, Any]] = None, device: Union[None, str, torch.device] = None, random_seed: Optional[int] = None, @@ -915,6 +916,7 @@ def pipeline( # noqa: C901 training_loop_instance: TrainingLoop = training_loop( model=model_instance, optimizer=optimizer_instance, + automatic_memory_optimization=automatic_memory_optimization, ) elif training_loop is not SLCWATrainingLoop: raise ValueError('Can not specify negative sampler with LCWA') @@ -927,14 +929,20 @@ def pipeline( # noqa: C901 training_loop_instance: TrainingLoop = SLCWATrainingLoop( model=model_instance, optimizer=optimizer_instance, + automatic_memory_optimization=automatic_memory_optimization, negative_sampler_cls=negative_sampler, negative_sampler_kwargs=negative_sampler_kwargs, ) evaluator = get_evaluator_cls(evaluator) - evaluator_instance: Evaluator = evaluator( - **(evaluator_kwargs or {}), - ) + # TODO @mehdi is setting the automatic memory optimization as an attribute + # of the class appropriate, since it doesn't cause any state to be stored? + # I think it might be better to have this as an argument to the + # Evaluator.evaluate() function instead + if evaluation_kwargs is None: + evaluator_kwargs = {} + evaluator_kwargs.setdefault('automatic_memory_optimization', automatic_memory_optimization) + evaluator_instance: Evaluator = evaluator(**evaluator_kwargs) if evaluation_kwargs is None: evaluation_kwargs = {} diff --git a/tests/test_models.py b/tests/test_models.py index df109de66c..5ca0c99b03 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -327,6 +327,16 @@ def cli_extras(self): for k, v in kwargs.items(): extras.append('--' + k.replace('_', '-')) extras.append(str(v)) + + # For the high/low memory test cases of NTN, SE, etc. + if self.training_loop_kwargs and 'automatic_memory_optimization' in self.training_loop_kwargs: + automatic_memory_optimization = self.training_loop_kwargs.get('automatic_memory_optimization') + if automatic_memory_optimization is True: + extras.append('--automatic-memory-optimization') + elif automatic_memory_optimization is False: + extras.append('--no-automatic-memory-optimization') + # else, leave to default + extras += [ '--number-epochs', self.train_num_epochs, '--embedding-dim', self.embedding_dim,