🔥⚡ POC PyTorch Lightning integration (pykeen#905)

Co-authored-by: Charles Tapley Hoyt <[email protected]>

Author: mberr

docs/source/contrib/lightning.rst

@@ -0,0 +1,6 @@
+PyTorch Lightning Integration
+=============================
+.. automodapi:: pykeen.contrib.lightning
+    :no-heading:
+    :headings: --
+    :no-inheritance-diagram:

docs/source/index.rst

@@ -20,6 +20,7 @@ PyKEEN
    tutorial/representations
    tutorial/node_piece
    tutorial/inductive_lp
+   contrib/lightning
 
 .. toctree::
    :caption: Bring Your Own

setup.cfg

@@ -101,6 +101,8 @@ tensorboard =
     tensorboard
 transformers =
     transformers
+lightning =
+    pytorch_lightning
 tests =
     unittest-templates>=0.0.5
     coverage

src/pykeen/contrib/__init__.py

@@ -0,0 +1,3 @@
+# -*- coding: utf-8 -*-
+
+"""This module contains code for non-standard dependencies in PyKEEN."""

src/pykeen/contrib/lightning.py

@@ -0,0 +1,165 @@
+# -*- coding: utf-8 -*-
+
+"""PyTorch Lightning integration.
+
+PyTorch Lightning poses an alternative way to implement a training
+loop and evaluation loop for knowledge graph embedding models that
+has some nice features:
+
+- mixed precision training
+- multi-gpu training
+
+.. code-block:: python
+
+    model = LitLCWAModule(
+        dataset="fb15k237",
+        dataset_kwargs=dict(create_inverse_triples=True),
+        model="mure",
+        model_kwargs=dict(embedding_dim=128, loss="bcewithlogits"),
+        batch_size=128,
+    )
+    trainer = pytorch_lightning.Trainer(
+        accelerator="auto",  # automatically choose accelerator
+        logger=False,  # defaults to TensorBoard; explicitly disabled here
+        precision=16,  # mixed precision training
+    )
+    trainer.fit(model=model)
+
+"""
+
+import pytorch_lightning
+import torch
+import torch.utils.data
+from class_resolver import HintOrType, OptionalKwargs
+
+from pykeen.datasets import get_dataset
+from pykeen.datasets.base import Dataset
+from pykeen.models import Model, model_resolver
+from pykeen.optimizers import optimizer_resolver
+from pykeen.triples.triples_factory import CoreTriplesFactory
+
+__all__ = [
+    "LitLCWAModule",
+]
+
+
+class LitLCWAModule(pytorch_lightning.LightningModule):
+    """A PyTorch Lightning module for training a model with LCWA training loop.
+
+    .. seealso:: https://github.com/pykeen/pykeen/pull/905
+    """
+
+    def __init__(
+        self,
+        # dataset
+        dataset: HintOrType[Dataset] = "nations",
+        dataset_kwargs: OptionalKwargs = None,
+        # model
+        model: HintOrType[Model] = "distmult",
+        model_kwargs: OptionalKwargs = None,
+        # stored outside of the training loop / optimizer to give access to auto-tuning from Lightning
+        batch_size: int = 32,
+        learning_rate: float = 1.0e-03,
+        # optimizer
+        optimizer: HintOrType[torch.optim.Optimizer] = None,
+        optimizer_kwargs: OptionalKwargs = None,
+    ):
+        """
+        Create the lightning module.
+
+        :param dataset:
+            the dataset, or a hint thereof
+        :param dataset_kwargs:
+            additional keyword-based parameters passed to the dataset
+
+        :param model:
+            the model, or a hint thereof
+        :param model_kwargs:
+            additional keyword-based parameters passed to the model
+
+        :param batch_size:
+            the training batch size
+        :param learning_rate:
+            the learning rate
+
+        :param optimizer:
+            the optimizer, or a hint thereof
+        :param optimizer_kwargs:
+            additional keyword-based parameters passed to the optimizer. should not contain `lr`, or `params`.
+        """
+        super().__init__()
+        self.dataset = get_dataset(dataset=dataset, dataset_kwargs=dataset_kwargs)
+        self.model = model_resolver.make(model, model_kwargs, triples_factory=self.dataset.training)
+        self.loss = self.model.loss
+        self.optimizer = optimizer
+        self.optimizer_kwargs = optimizer_kwargs
+        self.learning_rate = learning_rate
+        self.batch_size = batch_size
+
+    def forward(self, x):
+        """
+        Perform the prediction or inference step.
+
+        .. note::
+            in lightning, forward defines the prediction/inference actions
+        """
+        return self.model.predict_t(x)
+
+    def _step(self, batch, prefix: str):
+        """Refactored step."""
+        hr_batch, labels = batch
+        scores = self.model.score_t(hr_batch=hr_batch)
+        loss = self.loss.process_lcwa_scores(predictions=scores, labels=labels)
+        self.log(f"{prefix}_loss", loss)
+        return loss
+
+    def training_step(self, batch, batch_idx):
+        """Perform a training step."""
+        return self._step(batch, prefix="train")
+
+    def validation_step(self, batch, batch_idx, *args, **kwargs):
+        """Perform a validation step."""
+        return self._step(batch, prefix="val")
+
+    def _dataloader(self, triples_factory: CoreTriplesFactory, shuffle: bool = False) -> torch.utils.data.DataLoader:
+        """Create a data loader."""
+        return torch.utils.data.DataLoader(
+            dataset=triples_factory.create_lcwa_instances(),
+            batch_size=self.batch_size,
+            shuffle=shuffle,
+        )
+
+    def train_dataloader(self):
+        """Create the training data loader."""
+        return self._dataloader(triples_factory=self.dataset.training, shuffle=True)
+
+    def val_dataloader(self):
+        """Create the validation data loader."""
+        return self._dataloader(triples_factory=self.dataset.validation, shuffle=False)
+
+    def configure_optimizers(self):
+        """Configure the optimizers."""
+        return optimizer_resolver.make(
+            self.optimizer, self.optimizer_kwargs, params=self.parameters(), lr=self.learning_rate
+        )
+
+
+def _main():
+    """Run PyTorch lightning model."""
+    model = LitLCWAModule(
+        dataset="fb15k237",
+        dataset_kwargs=dict(create_inverse_triples=True),
+        model="mure",
+        model_kwargs=dict(embedding_dim=128, loss="bcewithlogits"),
+        batch_size=128,
+    )
+    trainer = pytorch_lightning.Trainer(
+        accelerator="auto",  # automatically choose accelerator
+        logger=False,  # defaults to TensorBoard; explicitly disabled here
+        precision=16,  # mixed precision training
+    )
+    trainer.fit(model=model)
+
+
+if __name__ == "__main__":
+    _main()

tox.ini

@@ -190,6 +190,7 @@ description = Test building the documentation in an isolated environment.
 changedir = docs
 extras =
     docs
+    lightning
 commands =
     mkdir -p {envtmpdir}
     cp -r source {envtmpdir}/source
@@ -209,7 +210,7 @@ whitelist_externals =
 [testenv:docs]
 description = Build the documentation locally.
 extras =
-    docs
+    {[testenv:docs-test]extras}
 commands =
     python -m sphinx -W -b html -d docs/build/doctrees docs/source docs/build/html