From 9191dbbbdecfc2f550671e6dbc6b7ef87a9f9515 Mon Sep 17 00:00:00 2001 From: Mehdi Ali <33023925+mali-git@users.noreply.github.com> Date: Fri, 27 Nov 2020 21:26:47 +0100 Subject: [PATCH] Remove Automatic Memory Optimization Attribute From Model (#176) * Pass automatic memory optimization flag to train and evaluation fcts. * Remove automatic memory optimization flag from model * Remove automatic memory optimization flag from model * Pass missing argument, extend docstring * Adapt docstring * Adapt automatic memory optimization flag * Revert * Define automatic memory optimization as an attribute of Evaluator * Define automatic memory optimization as an attribute * Adapt test to new API * Remove automatic memory optimization flag from model * Pass flake 8 * Pass tests * Refactor * Pass automatic memory optimization to super class * Fix low/high memory tests * Adapt early stopping tests * Adapt evaluator tests * Adapt training tests * Make passing explicit and fix type annotations * Update test_early_stopping.py * Update test_evaluators.py * Enable use of amo in pipeline and CLI * Fix typo Co-authored-by: Charles Tapley Hoyt --- src/pykeen/evaluation/evaluator.py | 8 +++-- src/pykeen/evaluation/rank_based_evaluator.py | 3 +- src/pykeen/evaluation/sklearn.py | 9 +++-- src/pykeen/models/base.py | 15 -------- src/pykeen/models/cli/builders.py | 3 ++ src/pykeen/models/cli/options.py | 5 +++ .../models/multimodal/complex_literal.py | 2 -- .../models/multimodal/distmult_literal.py | 2 -- src/pykeen/models/unimodal/complex.py | 5 --- src/pykeen/models/unimodal/conv_e.py | 2 -- src/pykeen/models/unimodal/conv_kb.py | 2 -- src/pykeen/models/unimodal/distmult.py | 2 -- src/pykeen/models/unimodal/ermlp.py | 2 -- src/pykeen/models/unimodal/ermlpe.py | 2 -- src/pykeen/models/unimodal/hole.py | 2 -- src/pykeen/models/unimodal/kg2e.py | 2 -- src/pykeen/models/unimodal/ntn.py | 2 -- src/pykeen/models/unimodal/proj_e.py | 2 -- src/pykeen/models/unimodal/rescal.py | 2 -- src/pykeen/models/unimodal/rgcn.py | 2 -- src/pykeen/models/unimodal/rotate.py | 2 -- src/pykeen/models/unimodal/simple.py | 2 -- .../models/unimodal/structured_embedding.py | 2 -- src/pykeen/models/unimodal/trans_d.py | 2 -- src/pykeen/models/unimodal/trans_e.py | 2 -- src/pykeen/models/unimodal/trans_h.py | 2 -- src/pykeen/models/unimodal/trans_r.py | 2 -- src/pykeen/models/unimodal/tucker.py | 2 -- .../models/unimodal/unstructured_model.py | 2 -- src/pykeen/pipeline.py | 14 ++++++-- src/pykeen/training/slcwa.py | 5 +++ src/pykeen/training/training_loop.py | 12 +++++-- tests/test_early_stopping.py | 17 ++++----- tests/test_evaluators.py | 12 +++---- tests/test_models.py | 33 +++++++++++++---- tests/test_training.py | 35 +++++++++++++------ 36 files changed, 115 insertions(+), 105 deletions(-) diff --git a/src/pykeen/evaluation/evaluator.py b/src/pykeen/evaluation/evaluator.py index 6af4860259..ea3a1af57a 100644 --- a/src/pykeen/evaluation/evaluator.py +++ b/src/pykeen/evaluation/evaluator.py @@ -65,13 +65,15 @@ def __init__( self, filtered: bool = False, requires_positive_mask: bool = False, - batch_size: int = None, - slice_size: int = None, + batch_size: Optional[int] = None, + slice_size: Optional[int] = None, + automatic_memory_optimization: bool = True, ): self.filtered = filtered self.requires_positive_mask = requires_positive_mask self.batch_size = batch_size self.slice_size = slice_size + self.automatic_memory_optimization = automatic_memory_optimization @classmethod def get_normalized_name(cls) -> str: @@ -135,7 +137,7 @@ def evaluate( if mapped_triples is None: mapped_triples = model.triples_factory.mapped_triples - if batch_size is None and model.automatic_memory_optimization: + if batch_size is None and self.automatic_memory_optimization: batch_size, slice_size = self.batch_and_slice( model=model, mapped_triples=mapped_triples, diff --git a/src/pykeen/evaluation/rank_based_evaluator.py b/src/pykeen/evaluation/rank_based_evaluator.py index 0093d7ca6c..d4e51fa21b 100644 --- a/src/pykeen/evaluation/rank_based_evaluator.py +++ b/src/pykeen/evaluation/rank_based_evaluator.py @@ -218,6 +218,7 @@ def __init__( self, ks: Optional[Iterable[Union[int, float]]] = None, filtered: bool = True, + automatic_memory_optimization: bool = True, **kwargs, ): """Initialize rank-based evaluator. @@ -228,7 +229,7 @@ def __init__( Whether to use the filtered evaluation protocol. If enabled, ranking another true triple higher than the currently considered one will not decrease the score. """ - super().__init__(filtered=filtered, **kwargs) + super().__init__(filtered=filtered, automatic_memory_optimization=automatic_memory_optimization, **kwargs) self.ks = tuple(ks) if ks is not None else (1, 3, 5, 10) for k in self.ks: if isinstance(k, float) and not (0 < k < 1): diff --git a/src/pykeen/evaluation/sklearn.py b/src/pykeen/evaluation/sklearn.py index e93abb9694..0a65a5f1ce 100644 --- a/src/pykeen/evaluation/sklearn.py +++ b/src/pykeen/evaluation/sklearn.py @@ -68,8 +68,13 @@ def get_metric(self, name: str) -> float: # noqa: D102 class SklearnEvaluator(Evaluator): """An evaluator that uses a Scikit-learn metric.""" - def __init__(self, **kwargs): - super().__init__(filtered=False, requires_positive_mask=True, **kwargs) + def __init__(self, automatic_memory_optimization: bool = True, **kwargs): + super().__init__( + filtered=False, + requires_positive_mask=True, + automatic_memory_optimization=automatic_memory_optimization, + **kwargs, + ) self.all_scores = {} self.all_positives = {} diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 26e933c8b0..3e3fa6302f 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -231,7 +231,6 @@ def __init__( triples_factory: TriplesFactory, loss: Optional[Loss] = None, predict_with_sigmoid: bool = False, - automatic_memory_optimization: Optional[bool] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, @@ -246,10 +245,6 @@ def __init__( Whether to apply sigmoid onto the scores when predicting scores. Applying sigmoid at prediction time may lead to exactly equal scores for certain triples with very high, or very low score. When not trained with applying sigmoid (or using BCEWithLogitsLoss), the scores are not calibrated to perform well with sigmoid. - :param automatic_memory_optimization: - If set to `True`, the model derives the maximum possible batch sizes for the scoring of triples during - evaluation and also training (if no batch size was given). This allows to fully utilize the hardware at hand - and achieves the fastest calculations possible. :param preferred_device: The preferred device for model training and inference. :param random_seed: @@ -268,9 +263,6 @@ def __init__( elif random_seed is not NoRandomSeedNecessary: set_random_seed(random_seed) - if automatic_memory_optimization is None: - automatic_memory_optimization = True - # Loss if loss is None: self.loss = self.loss_default(**self.loss_default_kwargs) @@ -299,9 +291,6 @@ def __init__( ''' self.predict_with_sigmoid = predict_with_sigmoid - # This allows to store the optimized parameters - self.automatic_memory_optimization = automatic_memory_optimization - @classmethod def _is_abstract(cls) -> bool: return inspect.isabstract(cls) @@ -1041,7 +1030,6 @@ def __init__( embedding_dim: int = 50, loss: Optional[Loss] = None, predict_with_sigmoid: bool = False, - automatic_memory_optimization: Optional[bool] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, @@ -1062,7 +1050,6 @@ def __init__( """ super().__init__( triples_factory=triples_factory, - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, @@ -1105,7 +1092,6 @@ def __init__( relation_dim: Optional[int] = None, loss: Optional[Loss] = None, predict_with_sigmoid: bool = False, - automatic_memory_optimization: Optional[bool] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, @@ -1133,7 +1119,6 @@ def __init__( """ super().__init__( triples_factory=triples_factory, - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, 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/models/multimodal/complex_literal.py b/src/pykeen/models/multimodal/complex_literal.py index 9e7d585ea2..a63daf21f2 100644 --- a/src/pykeen/models/multimodal/complex_literal.py +++ b/src/pykeen/models/multimodal/complex_literal.py @@ -37,7 +37,6 @@ def __init__( self, triples_factory: TriplesNumericLiteralsFactory, embedding_dim: int = 50, - automatic_memory_optimization: Optional[bool] = None, input_dropout: float = 0.2, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, @@ -47,7 +46,6 @@ def __init__( super().__init__( triples_factory=triples_factory, embedding_dim=embedding_dim, - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/multimodal/distmult_literal.py b/src/pykeen/models/multimodal/distmult_literal.py index 67090a0a5c..2838515534 100644 --- a/src/pykeen/models/multimodal/distmult_literal.py +++ b/src/pykeen/models/multimodal/distmult_literal.py @@ -31,7 +31,6 @@ def __init__( self, triples_factory: TriplesNumericLiteralsFactory, embedding_dim: int = 50, - automatic_memory_optimization: Optional[bool] = None, input_dropout: float = 0.0, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, @@ -40,7 +39,6 @@ def __init__( super().__init__( triples_factory=triples_factory, embedding_dim=embedding_dim, - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index 44834ec5e2..f205304c55 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -70,7 +70,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 200, - automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, @@ -82,9 +81,6 @@ def __init__( The triple factory connected to the model. :param embedding_dim: The embedding dimensionality of the entity embeddings. - :param automatic_memory_optimization: bool - Whether to automatically optimize the sub-batch size during training and batch size during evaluation with - regards to the hardware at hand. :param loss: OptionalLoss (optional) The loss to use. Defaults to SoftplusLoss. :param preferred_device: str (optional) @@ -97,7 +93,6 @@ def __init__( super().__init__( triples_factory=triples_factory, embedding_dim=2 * embedding_dim, # complex embeddings - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index bbe964ff84..c53f5c0ffa 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -177,7 +177,6 @@ def __init__( output_dropout: float = 0.3, feature_map_dropout: float = 0.2, embedding_dim: int = 200, - automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, @@ -196,7 +195,6 @@ def __init__( super().__init__( triples_factory=triples_factory, embedding_dim=embedding_dim, - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/conv_kb.py b/src/pykeen/models/unimodal/conv_kb.py index 24673434b8..a8db30e1de 100644 --- a/src/pykeen/models/unimodal/conv_kb.py +++ b/src/pykeen/models/unimodal/conv_kb.py @@ -76,7 +76,6 @@ def __init__( triples_factory: TriplesFactory, hidden_dropout_rate: float = 0., embedding_dim: int = 200, - automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, num_filters: int = 400, @@ -91,7 +90,6 @@ def __init__( triples_factory=triples_factory, embedding_dim=embedding_dim, loss=loss, - automatic_memory_optimization=automatic_memory_optimization, preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index 455defd809..428d0f0c56 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -71,7 +71,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 50, - automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, @@ -84,7 +83,6 @@ def __init__( super().__init__( triples_factory=triples_factory, embedding_dim=embedding_dim, - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index 7fded52e84..86fbb32585 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -45,7 +45,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 50, - automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, @@ -56,7 +55,6 @@ def __init__( super().__init__( triples_factory=triples_factory, embedding_dim=embedding_dim, - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/ermlpe.py b/src/pykeen/models/unimodal/ermlpe.py index 05976b4bd6..8b4f73c771 100644 --- a/src/pykeen/models/unimodal/ermlpe.py +++ b/src/pykeen/models/unimodal/ermlpe.py @@ -60,7 +60,6 @@ def __init__( input_dropout: float = 0.2, hidden_dropout: float = 0.3, embedding_dim: int = 200, - automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, @@ -69,7 +68,6 @@ def __init__( super().__init__( triples_factory=triples_factory, embedding_dim=embedding_dim, - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/hole.py b/src/pykeen/models/unimodal/hole.py index 9dc5d21143..f4a5f04aac 100644 --- a/src/pykeen/models/unimodal/hole.py +++ b/src/pykeen/models/unimodal/hole.py @@ -57,7 +57,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 200, - automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, @@ -68,7 +67,6 @@ def __init__( triples_factory=triples_factory, embedding_dim=embedding_dim, loss=loss, - automatic_memory_optimization=automatic_memory_optimization, preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index 846f887acc..56d0f4ec52 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -61,7 +61,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 50, - automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, @@ -80,7 +79,6 @@ def __init__( super().__init__( triples_factory=triples_factory, embedding_dim=embedding_dim, - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/ntn.py b/src/pykeen/models/unimodal/ntn.py index 30c28d3b75..47ae282f38 100644 --- a/src/pykeen/models/unimodal/ntn.py +++ b/src/pykeen/models/unimodal/ntn.py @@ -55,7 +55,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 100, - automatic_memory_optimization: Optional[bool] = None, num_slices: int = 4, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, @@ -73,7 +72,6 @@ def __init__( super().__init__( triples_factory=triples_factory, embedding_dim=embedding_dim, - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/proj_e.py b/src/pykeen/models/unimodal/proj_e.py index 98e39e5a59..2890866ad9 100644 --- a/src/pykeen/models/unimodal/proj_e.py +++ b/src/pykeen/models/unimodal/proj_e.py @@ -59,7 +59,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 50, - automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, @@ -69,7 +68,6 @@ def __init__( super().__init__( triples_factory=triples_factory, embedding_dim=embedding_dim, - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/rescal.py b/src/pykeen/models/unimodal/rescal.py index d158b1e5cf..4f097dc5e6 100644 --- a/src/pykeen/models/unimodal/rescal.py +++ b/src/pykeen/models/unimodal/rescal.py @@ -54,7 +54,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 50, - automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, @@ -72,7 +71,6 @@ def __init__( triples_factory=triples_factory, embedding_dim=embedding_dim, relation_dim=embedding_dim ** 2, # d x d matrices - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/rgcn.py b/src/pykeen/models/unimodal/rgcn.py index 79bd5a61f7..025e45ee5b 100644 --- a/src/pykeen/models/unimodal/rgcn.py +++ b/src/pykeen/models/unimodal/rgcn.py @@ -486,7 +486,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 500, - automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, predict_with_sigmoid: bool = False, preferred_device: DeviceHint = None, @@ -511,7 +510,6 @@ def __init__( raise ValueError('R-GCN handles edges in an undirected manner.') super().__init__( triples_factory=triples_factory, - automatic_memory_optimization=automatic_memory_optimization, loss=loss, predict_with_sigmoid=predict_with_sigmoid, preferred_device=preferred_device, diff --git a/src/pykeen/models/unimodal/rotate.py b/src/pykeen/models/unimodal/rotate.py index e1d3501447..7eb1e007fc 100644 --- a/src/pykeen/models/unimodal/rotate.py +++ b/src/pykeen/models/unimodal/rotate.py @@ -85,7 +85,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 200, - automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, @@ -95,7 +94,6 @@ def __init__( triples_factory=triples_factory, embedding_dim=2 * embedding_dim, loss=loss, - automatic_memory_optimization=automatic_memory_optimization, preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index 9684334355..11ec58424e 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -66,7 +66,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 200, - automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, @@ -76,7 +75,6 @@ def __init__( super().__init__( triples_factory=triples_factory, embedding_dim=embedding_dim, - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/structured_embedding.py b/src/pykeen/models/unimodal/structured_embedding.py index 1401504172..eb2fbbd362 100644 --- a/src/pykeen/models/unimodal/structured_embedding.py +++ b/src/pykeen/models/unimodal/structured_embedding.py @@ -51,7 +51,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 50, - automatic_memory_optimization: Optional[bool] = None, scoring_fct_norm: int = 1, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, @@ -66,7 +65,6 @@ def __init__( super().__init__( triples_factory=triples_factory, embedding_dim=embedding_dim, - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index 6042d7dfa5..47d69627d4 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -110,7 +110,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 50, - automatic_memory_optimization: Optional[bool] = None, relation_dim: int = 30, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, @@ -121,7 +120,6 @@ def __init__( triples_factory=triples_factory, embedding_dim=embedding_dim, relation_dim=relation_dim, - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/trans_e.py b/src/pykeen/models/unimodal/trans_e.py index 0f4cfeecc2..e8c8afe193 100644 --- a/src/pykeen/models/unimodal/trans_e.py +++ b/src/pykeen/models/unimodal/trans_e.py @@ -51,7 +51,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 50, - automatic_memory_optimization: Optional[bool] = None, scoring_fct_norm: int = 1, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, @@ -70,7 +69,6 @@ def __init__( super().__init__( triples_factory=triples_factory, embedding_dim=embedding_dim, - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index b233615a82..cbf3853fe5 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -66,7 +66,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 50, - automatic_memory_optimization: Optional[bool] = None, scoring_fct_norm: int = 1, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, @@ -81,7 +80,6 @@ def __init__( super().__init__( triples_factory=triples_factory, embedding_dim=embedding_dim, - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index 473dffa8dd..760cb3541f 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -76,7 +76,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 50, - automatic_memory_optimization: Optional[bool] = None, relation_dim: int = 30, scoring_fct_norm: int = 1, loss: Optional[Loss] = None, @@ -89,7 +88,6 @@ def __init__( triples_factory=triples_factory, embedding_dim=embedding_dim, relation_dim=relation_dim, - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/tucker.py b/src/pykeen/models/unimodal/tucker.py index 8dd9af3cfd..0ac4afcb45 100644 --- a/src/pykeen/models/unimodal/tucker.py +++ b/src/pykeen/models/unimodal/tucker.py @@ -80,7 +80,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 200, - automatic_memory_optimization: Optional[bool] = None, relation_dim: Optional[int] = None, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, @@ -104,7 +103,6 @@ def __init__( triples_factory=triples_factory, embedding_dim=embedding_dim, relation_dim=relation_dim, - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/unstructured_model.py b/src/pykeen/models/unimodal/unstructured_model.py index ef3265c7c0..6370292e21 100644 --- a/src/pykeen/models/unimodal/unstructured_model.py +++ b/src/pykeen/models/unimodal/unstructured_model.py @@ -47,7 +47,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 50, - automatic_memory_optimization: Optional[bool] = None, scoring_fct_norm: int = 1, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, @@ -62,7 +61,6 @@ def __init__( super().__init__( triples_factory=triples_factory, embedding_dim=embedding_dim, - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/pipeline.py b/src/pykeen/pipeline.py index 472c9c76c8..37ab4550f2 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 evaluator_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/src/pykeen/training/slcwa.py b/src/pykeen/training/slcwa.py index e3b9727359..e878a1688e 100644 --- a/src/pykeen/training/slcwa.py +++ b/src/pykeen/training/slcwa.py @@ -35,6 +35,7 @@ def __init__( optimizer: Optional[Optimizer] = None, negative_sampler_cls: Optional[Type[NegativeSampler]] = None, negative_sampler_kwargs: Optional[Mapping[str, Any]] = None, + automatic_memory_optimization: bool = True, ): """Initialize the training loop. @@ -43,10 +44,14 @@ def __init__( :param negative_sampler_cls: The class of the negative sampler :param negative_sampler_kwargs: Keyword arguments to pass to the negative sampler class on instantiation for every positive one + :param automatic_memory_optimization: + Whether to automatically optimize the sub-batch size during + training and batch size during evaluation with regards to the hardware at hand. """ super().__init__( model=model, optimizer=optimizer, + automatic_memory_optimization=automatic_memory_optimization, ) if negative_sampler_cls is None: diff --git a/src/pykeen/training/training_loop.py b/src/pykeen/training/training_loop.py index 2048193514..ab5d2cca66 100644 --- a/src/pykeen/training/training_loop.py +++ b/src/pykeen/training/training_loop.py @@ -79,16 +79,21 @@ def __init__( self, model: Model, optimizer: Optional[Optimizer] = None, + automatic_memory_optimization: bool = True, ) -> None: """Initialize the training loop. :param model: The model to train :param optimizer: The optimizer to use while training the model + :param automatic_memory_optimization: bool + Whether to automatically optimize the sub-batch size during + training and batch size during evaluation with regards to the hardware at hand. """ self.model = model self.optimizer = optimizer self.training_instances = None self.losses_per_epochs = [] + self.automatic_memory_optimization = automatic_memory_optimization if self.loss_blacklist and isinstance(self.model.loss, tuple(self.loss_blacklist)): raise TrainingApproachLossMismatchError( @@ -146,6 +151,9 @@ def train( :param slice_size: >0 The divisor for the scoring function when using slicing. This is only possible for LCWA training loops in general and only for models that have the slicing capability implemented. + :param automatic_memory_optimization: bool + Whether to automatically optimize the sub-batch size during training and batch size during evaluation with + regards to the hardware at hand. :param label_smoothing: (0 <= label_smoothing < 1) If larger than zero, use label smoothing. :param sampler: (None or 'schlichtkrull') @@ -264,13 +272,13 @@ def _train( # noqa: C901 # Take the biggest possible training batch_size, if batch_size not set batch_size_sufficient = False if batch_size is None: - if self.model.automatic_memory_optimization: + if self.automatic_memory_optimization: batch_size, batch_size_sufficient = self.batch_size_search() else: batch_size = 256 # This will find necessary parameters to optimize the use of the hardware at hand - if not only_size_probing and self.model.automatic_memory_optimization and not batch_size_sufficient: + if not only_size_probing and self.automatic_memory_optimization and not batch_size_sufficient: # return the relevant parameters slice_size and batch_size sub_batch_size, slice_size = self.sub_batch_and_slice(batch_size) diff --git a/tests/test_early_stopping.py b/tests/test_early_stopping.py index dc4946371d..907a4e627a 100644 --- a/tests/test_early_stopping.py +++ b/tests/test_early_stopping.py @@ -44,8 +44,8 @@ def test_is_improvement(): class MockEvaluator(Evaluator): """A mock evaluator for testing early stopping.""" - def __init__(self, losses: Iterable[float]) -> None: - super().__init__() + def __init__(self, losses: Iterable[float], automatic_memory_optimization: bool = True) -> None: + super().__init__(automatic_memory_optimization=automatic_memory_optimization) self.losses = tuple(losses) self.losses_iter = iter(self.losses) @@ -105,8 +105,8 @@ def __repr__(self): # noqa: D105 class MockModel(EntityRelationEmbeddingModel): """A mock model returning fake scores.""" - def __init__(self, triples_factory: TriplesFactory, automatic_memory_optimization: bool): - super().__init__(triples_factory=triples_factory, automatic_memory_optimization=automatic_memory_optimization) + def __init__(self, triples_factory: TriplesFactory): + super().__init__(triples_factory=triples_factory) num_entities = self.num_entities self.scores = torch.arange(num_entities, dtype=torch.float) @@ -167,10 +167,10 @@ class TestEarlyStopping(unittest.TestCase): def setUp(self): """Prepare for testing the early stopper.""" - self.mock_evaluator = MockEvaluator(self.mock_losses) # Set automatic_memory_optimization to false for tests + self.mock_evaluator = MockEvaluator(self.mock_losses, automatic_memory_optimization=False) nations = Nations() - self.model = MockModel(triples_factory=nations.training, automatic_memory_optimization=False) + self.model = MockModel(triples_factory=nations.training) self.stopper = EarlyStopper( model=self.model, evaluator=self.mock_evaluator, @@ -253,8 +253,8 @@ def test_early_stopping(self): """Tests early stopping.""" # Set automatic_memory_optimization to false during testing nations = Nations() - model: Model = TransE(triples_factory=nations.training, automatic_memory_optimization=False) - evaluator = RankBasedEvaluator() + model: Model = TransE(triples_factory=nations.training) + evaluator = RankBasedEvaluator(automatic_memory_optimization=False) stopper = EarlyStopper( model=model, evaluator=evaluator, @@ -266,6 +266,7 @@ def test_early_stopping(self): training_loop = SLCWATrainingLoop( model=model, optimizer=Adam(params=model.get_grad_params()), + automatic_memory_optimization=False, ) losses = training_loop.train( num_epochs=self.max_num_epochs, diff --git a/tests/test_evaluators.py b/tests/test_evaluators.py index a7068d4827..ad9c216326 100644 --- a/tests/test_evaluators.py +++ b/tests/test_evaluators.py @@ -394,8 +394,8 @@ def test_filter_corrupted_triples(self): class DummyEvaluator(Evaluator): """A dummy evaluator for testing the structure of the evaluation function.""" - def __init__(self, *, counter: int, filtered: bool) -> None: - super().__init__(filtered=filtered) + def __init__(self, *, counter: int, filtered: bool, automatic_memory_optimization: bool = True) -> None: + super().__init__(filtered=filtered, automatic_memory_optimization=automatic_memory_optimization) self.counter = counter def process_tail_scores_( @@ -431,8 +431,8 @@ def __repr__(self): # noqa: D105 class DummyModel(EntityRelationEmbeddingModel): """A dummy model returning fake scores.""" - def __init__(self, triples_factory: TriplesFactory, automatic_memory_optimization: bool): - super().__init__(triples_factory=triples_factory, automatic_memory_optimization=automatic_memory_optimization) + def __init__(self, triples_factory: TriplesFactory): + super().__init__(triples_factory=triples_factory) num_entities = self.num_entities self.scores = torch.arange(num_entities, dtype=torch.float) @@ -462,9 +462,9 @@ class TestEvaluationStructure(unittest.TestCase): def setUp(self): """Prepare for testing the evaluation structure.""" self.counter = 1337 - self.evaluator = DummyEvaluator(counter=self.counter, filtered=True) + self.evaluator = DummyEvaluator(counter=self.counter, filtered=True, automatic_memory_optimization=False) self.triples_factory = Nations().training - self.model = DummyModel(triples_factory=self.triples_factory, automatic_memory_optimization=False) + self.model = DummyModel(triples_factory=self.triples_factory) def test_evaluation_structure(self): """Test if the evaluator has a balanced call of head and tail processors.""" diff --git a/tests/test_models.py b/tests/test_models.py index 1064c6bece..5ca0c99b03 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -92,6 +92,9 @@ class _ModelTestCase: #: Additional arguments passed to the model's constructor method model_kwargs: ClassVar[Optional[Mapping[str, Any]]] = None + #: Additional arguments passed to the training loop's constructor method + training_loop_kwargs: ClassVar[Optional[Mapping[str, Any]]] = None + #: The triples factory instance factory: TriplesFactory @@ -244,6 +247,7 @@ def test_train_slcwa(self) -> None: loop = SLCWATrainingLoop( model=self.model, optimizer=Adagrad(params=self.model.get_grad_params(), lr=0.001), + **(self.training_loop_kwargs or {}), ) losses = self._safe_train_loop( loop, @@ -259,6 +263,7 @@ def test_train_lcwa(self) -> None: loop = LCWATrainingLoop( model=self.model, optimizer=Adagrad(params=self.model.get_grad_params(), lr=0.001), + **(self.training_loop_kwargs or {}), ) losses = self._safe_train_loop( loop, @@ -322,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, @@ -710,6 +725,9 @@ class TestNTNLowMemory(_BaseNTNTest): model_kwargs = { 'num_slices': 2, + } + + training_loop_kwargs = { 'automatic_memory_optimization': True, } @@ -719,6 +737,9 @@ class TestNTNHighMemory(_BaseNTNTest): model_kwargs = { 'num_slices': 2, + } + + training_loop_kwargs = { 'automatic_memory_optimization': False, } @@ -815,17 +836,17 @@ def _check_constraints(self): class TestSELowMemory(_BaseTestSE): """Tests SE with low memory.""" - model_kwargs = dict( - automatic_memory_optimization=True, - ) + training_loop_kwargs = { + 'automatic_memory_optimization': True, + } class TestSEHighMemory(_BaseTestSE): """Tests SE with low memory.""" - model_kwargs = dict( - automatic_memory_optimization=False, - ) + training_loop_kwargs = { + 'automatic_memory_optimization': False, + } class TestTransD(_DistanceModelTestCase, unittest.TestCase): diff --git a/tests/test_training.py b/tests/test_training.py index 179148ea16..0a2186669a 100644 --- a/tests/test_training.py +++ b/tests/test_training.py @@ -20,8 +20,12 @@ class DummyTrainingLoop(SLCWATrainingLoop): """A wrapper around SLCWATrainingLoop.""" - def __init__(self, model: Model, sub_batch_size: int): - super().__init__(model=model, optimizer=optim.Adam(lr=1.0, params=model.parameters())) + def __init__(self, model: Model, sub_batch_size: int, automatic_memory_optimization: bool = False): + super().__init__( + model=model, + optimizer=optim.Adam(lr=1.0, params=model.parameters()), + automatic_memory_optimization=automatic_memory_optimization, + ) self.sub_batch_size = sub_batch_size def _process_batch( @@ -50,8 +54,12 @@ def _process_batch( class NaNTrainingLoop(SLCWATrainingLoop): """A wrapper around SLCWATrainingLoop returning NaN losses.""" - def __init__(self, model: Model, patience: int): - super().__init__(model=model, optimizer=optim.Adam(lr=1.0, params=model.parameters())) + def __init__(self, model: Model, patience: int, automatic_memory_optimization: bool = False): + super().__init__( + model=model, + optimizer=optim.Adam(lr=1.0, params=model.parameters()), + automatic_memory_optimization=automatic_memory_optimization, + ) self.patience = patience def _process_batch( @@ -89,14 +97,22 @@ def setUp(self) -> None: def test_sub_batching(self): """Test if sub-batching works as expected.""" - model = TransE(triples_factory=self.triples_factory, automatic_memory_optimization=False) - training_loop = DummyTrainingLoop(model=model, sub_batch_size=self.sub_batch_size) + model = TransE(triples_factory=self.triples_factory) + training_loop = DummyTrainingLoop( + model=model, + sub_batch_size=self.sub_batch_size, + automatic_memory_optimization=False, + ) training_loop.train(num_epochs=1, batch_size=self.batch_size, sub_batch_size=self.sub_batch_size) def test_sub_batching_support(self): """Test if sub-batching works as expected.""" - model = ConvE(triples_factory=self.triples_factory, automatic_memory_optimization=False) - training_loop = DummyTrainingLoop(model=model, sub_batch_size=self.sub_batch_size) + model = ConvE(triples_factory=self.triples_factory) + training_loop = DummyTrainingLoop( + model=model, + sub_batch_size=self.sub_batch_size, + automatic_memory_optimization=False, + ) def _try_train(): """Call train method.""" @@ -117,7 +133,6 @@ def test_blacklist_loss_on_slcwa(self): model = TransE( triples_factory=self.triples_factory, loss=CrossEntropyLoss(), - automatic_memory_optimization=False, ) with self.assertRaises(TrainingApproachLossMismatchError): - NaNTrainingLoop(model=model, patience=2) + NaNTrainingLoop(model=model, patience=2, automatic_memory_optimization=False)