🚅⚡ Test Training with PyTorch Lightning

[mberr]

This PR adds tests for the PyTorch Lightning integration.

Issues

• post_parameter_update is not applied, cf. failing pipeline for KG2E; fixed in 9052412
• RuntimeError for CompGCN, RESCAL and TransH

Related

• Add PyKEEN to ecosystem-ci Lightning-AI/ecosystem-ci#50

[cthoyt]

todo: add [![Lightning](https://img.shields.io/badge/-Lightning-792ee5?logo=pytorchlightning&logoColor=white)](https://pytorchlightning.ai) to the readme

[mberr]

There are some RuntimeErrors for CompGCN, RESCAL, and TransH which do not occur for not training with PyL.

src/pykeen/contrib/lightning.py:288: in lit_pipeline
    trainer.fit(model=lit)
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:768: in fit
    self._call_and_handle_interrupt(
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:721: in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:809: in _fit_impl
    results = self._run(model, ckpt_path=self.ckpt_path)
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:1234: in _run
    results = self._run_stage()
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:1321: in _run_stage
    return self._run_train()
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:1351: in _run_train
    self.fit_loop.run()
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/loops/base.py:204: in run
    self.advance(*args, **kwargs)
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/loops/fit_loop.py:269: in advance
    self._outputs = self.epoch_loop.run(self._data_fetcher)
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/loops/base.py:204: in run
    self.advance(*args, **kwargs)
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py:208: in advance
    batch_output = self.batch_loop.run(batch, batch_idx)
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/loops/base.py:204: in run
    self.advance(*args, **kwargs)
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/loops/batch/training_batch_loop.py:88: in advance
    outputs = self.optimizer_loop.run(split_batch, optimizers, batch_idx)
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/loops/base.py:204: in run
    self.advance(*args, **kwargs)
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py:203: in advance
    result = self._run_optimization(
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py:256: in _run_optimization
    self._optimizer_step(optimizer, opt_idx, batch_idx, closure)
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py:369: in _optimizer_step
    self.trainer._call_lightning_module_hook(
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:1593: in _call_lightning_module_hook
    output = fn(*args, **kwargs)
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/core/lightning.py:1625: in optimizer_step
    optimizer.step(closure=optimizer_closure)
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/core/optimizer.py:168: in step
    step_output = self._strategy.optimizer_step(self._optimizer, self._optimizer_idx, closure, **kwargs)
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/strategies/strategy.py:193: in optimizer_step
    return self.precision_plugin.optimizer_step(model, optimizer, opt_idx, closure, **kwargs)
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py:155: in optimizer_step
    return optimizer.step(closure=closure, **kwargs)
venv/venv-gpu/lib/python3.8/site-packages/torch/optim/optimizer.py:88: in wrapper
    return func(*args, **kwargs)
venv/venv-gpu/lib/python3.8/site-packages/torch/autograd/grad_mode.py:27: in decorate_context
    return func(*args, **kwargs)
venv/venv-gpu/lib/python3.8/site-packages/torch/optim/adam.py:100: in step
    loss = closure()
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py:140: in _wrap_closure
    closure_result = closure()
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py:148: in __call__
    self._result = self.closure(*args, **kwargs)
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py:143: in closure
    self._backward_fn(step_output.closure_loss)
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py:311: in backward_fn
    self.trainer._call_strategy_hook("backward", loss, optimizer, opt_idx)
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:1763: in _call_strategy_hook
    output = fn(*args, **kwargs)
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/strategies/strategy.py:168: in backward
    self.precision_plugin.backward(self.lightning_module, closure_loss, *args, **kwargs)
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py:80: in backward
    model.backward(closure_loss, optimizer, *args, **kwargs)
venv/venv-gpu/lib/python3.8/site-packages/pytorch_lightning/core/lightning.py:1370: in backward
    loss.backward(*args, **kwargs)
venv/venv-gpu/lib/python3.8/site-packages/torch/_tensor.py:363: in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)

[cthoyt]

what's the point of this? it doesn't return anything

[mberr]

It modifies the model's weights in-place

[cthoyt]

maybe add a bit of extra description then? for poor fools like me ;)

[mberr]

ca2f266

[cthoyt]

which ones doesn't it work for?

[mberr]

1d00631 - this comment was outdated.

[cthoyt]

@PyKEEN-bot test