Using Lightning DataModules

[simon-forb]

I have the following minimal code example:

from lightning import LightningDataModule
from torch.utils.data import random_split
from torch_geometric.data.lightning import LightningDataset
from torch_geometric.datasets import TUDataset
import lightning as L

trainer = L.Trainer()
reddit_dataset = TUDataset("./data", "REDDIT-BINARY")
train_set, val_set, test_set = random_split(reddit_dataset, [.7, .1, .2])
dm = LightningDataset(train_set, val_set, test_set, batch_size=32)

trainer.fit(model, datamodule=dm)

Which gives me the following error:

TypeError: An invalid dataloader was passed to `Trainer.fit(train_dataloaders=...)`. Found LightningDataset(train_dataset=<torch.utils.data.dataset.Subset object at 0x773852bfe710>, val_dataset=<torch.utils.data.dataset.Subset object at 0x773852c35110>, test_dataset=<torch.utils.data.dataset.Subset object at 0x773852bffa50>, batch_size=32, num_workers=0, pin_memory=True, persistent_workers=False).

I would have expected this is how LightningDataset is supposed to be applied.

My code works when I create standard DataLoaders from PyG.

Version: pyg 2.5.2 py311_torch_2.2.0_cu121

[rusty1s]

Which PL version are you using?

[simon-forb]

This is a very recent environment, it uses

pytorch-lightning=2.2.2=pyhd8ed1ab_0

[brenerrr]

I am having the same issue.

torch: 2.1.1
lightning: 2.4.0
torch_geometric: 2.4.0

[akihironitta]

Mind updating it to torch_geometric==2.6.0 and retrying it? I cannot reproduce this with the PyTorch Lightning version.

[zyjOwO]

I reimplemented similar issue while using Ray and PyG.
[Reason]
Ray (or your code) imports the [lightning] package, but PyG imports the [pytorch_lightning] package.
[Solution]
pyg.data.lightning.datamodule.py (line14)
from pytorch_lightning import LightningDataModule as PLLightningDataModule
==> from lightning import LightningDataModule as PLLightningDataModule
[Future]
[pytorch_lightning] got renamed to [lightning]. The newer PyG (>2.7.0) could consider it for avoiding such conflicts. @rusty1s

[akihironitta]

@zyjOwO Thanks for reporting this issue. Since lightning used to contain totally unnecessary code/dependencies until Lightning-AI/pytorch-lightning#20039, I'd prefer using pytorch-lightning, which the Lightning team has promised to keep supporting, but we could potentially consider adopting your suggested change. Feel free to send a PR to support both pytorch_lightning and lightning import names if you're interested in contributing to PyG!

[zyjOwO]

Thanks for your prompt response. I'd prefer using [lightning]. The Lightning team won't show [pytorch_lightning] anymore in new docs (lightning>=2.0.0). The Ray team used the following solution to support both packages. But it cannot avoid such conflicts totally.
[Solution]
try:
from lightning import *
except ModuleNotFoundError:
from pytorch_lightning import *