graphstorm
graphstorm copied to clipboard
feature: add support for PyTorch Lightning
Issue #, if available: N/A
Description of changes: Adds PyTorch Lightning support, LightningModule and LightningDataModule for Node GNN model and also a Jupyter nodebook demonstrating how it works.
By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.
Hello @clumsy took a closer look at the PR and had a couple of requests. We'd like to maintain a similar level of requirements between GSProcessing and this library.
- We should try to have some unit tests if possible.
- Let's put some more work pyproject.toml. Right now in the dev group I see a lot of dependencies that I'm not sure the project needs. For DGL in particular there's only support for py3.10 on Linux?
- Let's add some documentation. The example notebook is helpful but we also want to have a guide that helps users from installation to training, with AWS as the target cloud infra. At least one example on how to run this on SageMaker should be included.
- This also means that we likely have to include a compatible Dockerfile. If the existing SageMaker and local Dockerfiles can be slightly modified to add PTL support, we can have that in documentation.