import math from contextlib import contextmanager import pytest import torch import torch.nn.functional as F from torch import Tensor from torch_geometric.data import Data, HeteroData from torch_geometric.data.lightning import ( LightningDataset, LightningLinkData, LightningNodeData, ) from torch_geometric.nn import global_mean_pool from torch_geometric.sampler import BaseSampler, NeighborSampler from torch_geometric.testing import ( MyFeatureStore, MyGraphStore, get_random_edge_index, onlyCUDA, onlyFullTest, onlyNeighborSampler, onlyOnline, withPackage, ) try: from pytorch_lightning import LightningModule except ImportError: LightningModule = torch.nn.Module class LinearGraphModule(LightningModule): def __init__(self, in_channels: int, hidden_channels: int, out_channels: int): super().__init__() from torchmetrics import Accuracy self.lin1 = torch.nn.Linear(in_channels, hidden_channels) self.lin2 = torch.nn.Linear(hidden_channels, out_channels) self.train_acc = Accuracy(task='multiclass', num_classes=out_channels) self.val_acc = Accuracy(task='multiclass', num_classes=out_channels) self.test_acc = Accuracy(task='multiclass', num_classes=out_channels) def forward(self, x: Tensor, batch: Data) -> Tensor: # Basic test to ensure that the dataset is not replicated: self.trainer.datamodule.train_dataset._data.x.add_(1) x = self.lin1(x).relu() x = global_mean_pool(x, batch) x = self.lin2(x) return x def training_step(self, data: Data, batch_idx: int): y_hat = self(data.x, data.batch) loss = F.cross_entropy(y_hat, data.y) self.train_acc(y_hat.softmax(dim=-1), data.y) self.log('loss', loss, batch_size=data.num_graphs) self.log('train_acc', self.train_acc, batch_size=data.num_graphs) return loss def validation_step(self, data: Data, batch_idx: int): y_hat = self(data.x, data.batch) self.val_acc(y_hat.softmax(dim=-1), data.y) self.log('val_acc', self.val_acc, batch_size=data.num_graphs) def test_step(self, data: Data, batch_idx: int): y_hat = self(data.x, data.batch) self.test_acc(y_hat.softmax(dim=-1), data.y) self.log('test_acc', self.test_acc, batch_size=data.num_graphs) def configure_optimizers(self): return torch.optim.Adam(self.parameters(), lr=0.01) @onlyCUDA @onlyOnline @onlyFullTest @withPackage('pytorch_lightning>=2.0.0', 'torchmetrics>=0.11.0') @pytest.mark.parametrize('strategy_type', [None, 'ddp']) def test_lightning_dataset(get_dataset, strategy_type): import pytorch_lightning as pl from pytorch_lightning.utilities import rank_zero_only @contextmanager def expect_rank_zero_user_warning(match: str): if rank_zero_only.rank == 0: with pytest.warns(UserWarning, match=match): yield else: yield dataset = get_dataset(name='MUTAG').shuffle() train_dataset = dataset[:50] val_dataset = dataset[50:80] test_dataset = dataset[80:90] pred_dataset = dataset[90:] devices = 1 if strategy_type is None else torch.cuda.device_count() if strategy_type == 'ddp': strategy = pl.strategies.DDPStrategy(accelerator='gpu') else: strategy = pl.strategies.SingleDeviceStrategy(device='cuda:0') model = LinearGraphModule(dataset.num_features, 64, dataset.num_classes) trainer = pl.Trainer(strategy=strategy, devices=devices, max_epochs=1, log_every_n_steps=1) with pytest.warns(UserWarning, match="'shuffle=True' option is ignored"): datamodule = LightningDataset(train_dataset, val_dataset, test_dataset, pred_dataset, batch_size=5, num_workers=3, shuffle=True) assert 'shuffle' not in datamodule.kwargs old_x = train_dataset._data.x.clone() assert str(datamodule) == ('LightningDataset(train_dataset=MUTAG(50), ' 'val_dataset=MUTAG(30), ' 'test_dataset=MUTAG(10), ' 'pred_dataset=MUTAG(98), batch_size=5, ' 'num_workers=3, pin_memory=True, ' 'persistent_workers=True)') trainer.fit(model, datamodule) trainer.test(model, datamodule) new_x = train_dataset._data.x assert torch.all(new_x > old_x) # Ensure shared data. assert trainer.validate_loop._data_source.is_defined() assert trainer.test_loop._data_source.is_defined() # Test with `val_dataset=None` and `test_dataset=None`: if strategy_type is None: trainer = pl.Trainer(strategy=strategy, devices=devices, max_epochs=1, log_every_n_steps=1) datamodule = LightningDataset(train_dataset, batch_size=5) assert str(datamodule) == ('LightningDataset(train_dataset=MUTAG(50), ' 'batch_size=5, num_workers=0, ' 'pin_memory=True, ' 'persistent_workers=False)') with expect_rank_zero_user_warning("defined a `validation_step`"): trainer.fit(model, datamodule) assert not trainer.validate_loop._data_source.is_defined() assert not trainer.test_loop._data_source.is_defined() class LinearNodeModule(LightningModule): def __init__(self, in_channels: int, out_channels: int): super().__init__() from torchmetrics import Accuracy self.lin = torch.nn.Linear(in_channels, out_channels) self.train_acc = Accuracy(task='multiclass', num_classes=out_channels) self.val_acc = Accuracy(task='multiclass', num_classes=out_channels) self.test_acc = Accuracy(task='multiclass', num_classes=out_channels) def forward(self, x: Tensor) -> Tensor: # Basic test to ensure that the dataset is not replicated: self.trainer.datamodule.data.x.add_(1) return self.lin(x) def training_step(self, data: Data, batch_idx: int): y_hat = self(data.x)[data.train_mask] y = data.y[data.train_mask] loss = F.cross_entropy(y_hat, y) self.train_acc(y_hat.softmax(dim=-1), y) self.log('loss', loss, batch_size=y.size(0)) self.log('train_acc', self.train_acc, batch_size=y.size(0)) return loss def validation_step(self, data: Data, batch_idx: int): y_hat = self(data.x)[data.val_mask] y = data.y[data.val_mask] self.val_acc(y_hat.softmax(dim=-1), y) self.log('val_acc', self.val_acc, batch_size=y.size(0)) def test_step(self, data: Data, batch_idx: int): y_hat = self(data.x)[data.test_mask] y = data.y[data.test_mask] self.test_acc(y_hat.softmax(dim=-1), y) self.log('test_acc', self.test_acc, batch_size=y.size(0)) def configure_optimizers(self): return torch.optim.Adam(self.parameters(), lr=0.01) @onlyCUDA @onlyOnline @onlyFullTest @onlyNeighborSampler @withPackage('pytorch_lightning>=2.0.0', 'torchmetrics>=0.11.0', 'scipy') @pytest.mark.parametrize('loader', ['full', 'neighbor']) @pytest.mark.parametrize('strategy_type', [None, 'ddp']) def test_lightning_node_data(get_dataset, strategy_type, loader): import pytorch_lightning as pl dataset = get_dataset(name='Cora') data = dataset[0] data_repr = ('Data(x=[2708, 1433], edge_index=[2, 10556], y=[2708], ' 'train_mask=[2708], val_mask=[2708], test_mask=[2708])') model = LinearNodeModule(dataset.num_features, dataset.num_classes) if strategy_type is None or loader == 'full': devices = 1 else: devices = torch.cuda.device_count() if strategy_type == 'ddp': strategy = pl.strategies.DDPStrategy(accelerator='gpu') else: strategy = pl.strategies.SingleDeviceStrategy(device='cuda:0') if loader == 'full': # Set reasonable defaults for full-batch training: batch_size = 1 num_workers = 0 else: batch_size = 32 num_workers = 3 kwargs, kwargs_repr = {}, '' if loader == 'neighbor': kwargs['num_neighbors'] = [5] kwargs_repr += 'num_neighbors=[5], ' trainer = pl.Trainer(strategy=strategy, devices=devices, max_epochs=5, log_every_n_steps=1) datamodule = LightningNodeData(data, loader=loader, batch_size=batch_size, num_workers=num_workers, **kwargs) old_x = data.x.clone().cpu() assert str(datamodule) == (f'LightningNodeData(data={data_repr}, ' f'loader={loader}, batch_size={batch_size}, ' f'num_workers={num_workers}, {kwargs_repr}' f'pin_memory={loader != "full"}, ' f'persistent_workers={loader != "full"})') trainer.fit(model, datamodule) trainer.test(model, datamodule) new_x = data.x.cpu() assert torch.all(new_x > old_x) # Ensure shared data. assert trainer.validate_loop._data_source.is_defined() assert trainer.test_loop._data_source.is_defined() class LinearHeteroNodeModule(LightningModule): def __init__(self, in_channels: int, out_channels: int): super().__init__() from torchmetrics import Accuracy self.lin = torch.nn.Linear(in_channels, out_channels) self.train_acc = Accuracy(task='multiclass', num_classes=out_channels) self.val_acc = Accuracy(task='multiclass', num_classes=out_channels) self.test_acc = Accuracy(task='multiclass', num_classes=out_channels) def forward(self, x: Tensor) -> Tensor: # Basic test to ensure that the dataset is not replicated: self.trainer.datamodule.data['paper'].x.add_(1) return self.lin(x) def training_step(self, data: HeteroData, batch_idx: int): y_hat = self(data['paper'].x)[data['paper'].train_mask] y = data['paper'].y[data['paper'].train_mask] loss = F.cross_entropy(y_hat, y) self.train_acc(y_hat.softmax(dim=-1), y) self.log('loss', loss, batch_size=y.size(0)) self.log('train_acc', self.train_acc, batch_size=y.size(0)) return loss def validation_step(self, data: HeteroData, batch_idx: int): y_hat = self(data['paper'].x)[data['paper'].val_mask] y = data['paper'].y[data['paper'].val_mask] self.val_acc(y_hat.softmax(dim=-1), y) self.log('val_acc', self.val_acc, batch_size=y.size(0)) def test_step(self, data: HeteroData, batch_idx: int): y_hat = self(data['paper'].x)[data['paper'].test_mask] y = data['paper'].y[data['paper'].test_mask] self.test_acc(y_hat.softmax(dim=-1), y) self.log('test_acc', self.test_acc, batch_size=y.size(0)) def configure_optimizers(self): return torch.optim.Adam(self.parameters(), lr=0.01) @pytest.fixture def preserve_context(): num_threads = torch.get_num_threads() yield if torch.distributed.is_initialized(): torch.distributed.destroy_process_group() torch.set_num_threads(num_threads) @onlyCUDA @onlyFullTest @onlyNeighborSampler @withPackage('pytorch_lightning>=2.0.0', 'torchmetrics>=0.11.0') def test_lightning_hetero_node_data(preserve_context, get_dataset): import pytorch_lightning as pl data = get_dataset(name='hetero')[0] model = LinearHeteroNodeModule(data['paper'].num_features, int(data['paper'].y.max()) + 1) devices = torch.cuda.device_count() strategy = pl.strategies.DDPStrategy(accelerator='gpu') trainer = pl.Trainer(strategy=strategy, devices=devices, max_epochs=5, log_every_n_steps=1) datamodule = LightningNodeData(data, loader='neighbor', num_neighbors=[5], batch_size=32, num_workers=3) assert isinstance(datamodule.graph_sampler, NeighborSampler) original_x = data['paper'].x.clone() trainer.fit(model, datamodule) trainer.test(model, datamodule) assert torch.all(data['paper'].x > original_x) # Ensure shared data. assert trainer.validate_loop._data_source.is_defined() assert trainer.test_loop._data_source.is_defined() @withPackage('pytorch_lightning') def test_lightning_data_custom_sampler(): class DummySampler(BaseSampler): def sample_from_edges(self, *args, **kwargs): pass def sample_from_nodes(self, *args, **kwargs): pass data = Data(num_nodes=2, edge_index=torch.tensor([[0, 1], [1, 0]])) datamodule = LightningNodeData(data, node_sampler=DummySampler(), input_train_nodes=torch.arange(2)) assert isinstance(datamodule.graph_sampler, DummySampler) datamodule = LightningLinkData( data, link_sampler=DummySampler(), input_train_edges=torch.tensor([[0, 1], [0, 1]])) assert isinstance(datamodule.graph_sampler, DummySampler) @onlyCUDA @onlyFullTest @onlyNeighborSampler @withPackage('pytorch_lightning') def test_lightning_hetero_link_data(): torch.manual_seed(12345) data = HeteroData() data['paper'].x = torch.arange(10) data['author'].x = torch.arange(10) data['term'].x = torch.arange(10) data['paper', 'author'].edge_index = get_random_edge_index(10, 10, 10) data['author', 'paper'].edge_index = get_random_edge_index(10, 10, 10) data['paper', 'term'].edge_index = get_random_edge_index(10, 10, 10) data['author', 'term'].edge_index = get_random_edge_index(10, 10, 10) datamodule = LightningLinkData( data, input_train_edges=('author', 'paper'), input_val_edges=('paper', 'author'), input_test_edges=('paper', 'term'), input_pred_edges=('author', 'term'), loader='neighbor', num_neighbors=[5], batch_size=32, num_workers=0, ) assert isinstance(datamodule.graph_sampler, NeighborSampler) assert isinstance(datamodule.eval_graph_sampler, NeighborSampler) for batch in datamodule.train_dataloader(): assert 'edge_label_index' in batch['author', 'paper'] for batch in datamodule.val_dataloader(): assert 'edge_label_index' in batch['paper', 'author'] for batch in datamodule.test_dataloader(): assert 'edge_label_index' in batch['paper', 'term'] for batch in datamodule.predict_dataloader(): assert 'edge_label_index' in batch['author', 'term'] data['author'].time = torch.arange(data['author'].num_nodes) data['paper'].time = torch.arange(data['paper'].num_nodes) data['term'].time = torch.arange(data['term'].num_nodes) datamodule = LightningLinkData( data, input_train_edges=('author', 'paper'), input_train_time=torch.arange(data['author', 'paper'].num_edges), loader='neighbor', num_neighbors=[5], batch_size=32, num_workers=0, time_attr='time', ) for batch in datamodule.train_dataloader(): assert 'edge_label_index' in batch['author', 'paper'] assert 'edge_label_time' in batch['author', 'paper'] @onlyNeighborSampler @withPackage('pytorch_lightning') def test_lightning_hetero_link_data_custom_store(): torch.manual_seed(12345) feature_store = MyFeatureStore() graph_store = MyGraphStore() x = torch.arange(10) feature_store.put_tensor(x, group_name='paper', attr_name='x', index=None) feature_store.put_tensor(x, group_name='author', attr_name='x', index=None) feature_store.put_tensor(x, group_name='term', attr_name='x', index=None) edge_index = get_random_edge_index(10, 10, 10) graph_store.put_edge_index(edge_index=(edge_index[0], edge_index[1]), edge_type=('paper', 'to', 'author'), layout='coo', size=(10, 10)) graph_store.put_edge_index(edge_index=(edge_index[0], edge_index[1]), edge_type=('author', 'to', 'paper'), layout='coo', size=(10, 10)) graph_store.put_edge_index(edge_index=(edge_index[0], edge_index[1]), edge_type=('paper', 'to', 'term'), layout='coo', size=(10, 10)) datamodule = LightningLinkData( (feature_store, graph_store), input_train_edges=('author', 'to', 'paper'), loader='neighbor', num_neighbors=[5], batch_size=32, num_workers=0, ) batch = next(iter(datamodule.train_dataloader())) assert 'edge_label_index' in batch['author', 'paper'] @onlyOnline @onlyNeighborSampler @withPackage('pytorch_lightning', 'scipy') def test_eval_loader_kwargs(get_dataset): data = get_dataset(name='Cora')[0] node_sampler = NeighborSampler(data, num_neighbors=[5]) datamodule = LightningNodeData( data, node_sampler=node_sampler, batch_size=32, eval_loader_kwargs=dict(num_neighbors=[-1], batch_size=64), ) assert datamodule.loader_kwargs['batch_size'] == 32 assert datamodule.graph_sampler.num_neighbors.values == [5] assert datamodule.eval_loader_kwargs['batch_size'] == 64 assert datamodule.eval_graph_sampler.num_neighbors.values == [-1] train_loader = datamodule.train_dataloader() assert math.ceil(int(data.train_mask.sum()) / 32) == len(train_loader) val_loader = datamodule.val_dataloader() assert math.ceil(int(data.val_mask.sum()) / 64) == len(val_loader) test_loader = datamodule.test_dataloader() assert math.ceil(int(data.test_mask.sum()) / 64) == len(test_loader) pred_loader = datamodule.predict_dataloader() assert math.ceil(data.num_nodes / 64) == len(pred_loader)