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 <cthoyt@gmail.com>

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,

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):

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 = {}
 

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,

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:

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)),

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,

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,

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,

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,

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,

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,

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,

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,

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,

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,

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,

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,

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,