blog_4.py

class CoraDataset(LightningDataModule):

    NAME = "cora"
    
    ....

    def create_neighbor_sampler(self, batch_size=2, stage=None):
        # https://github.com/rusty1s/pytorch_geometric/tree/
        # master/torch_geometric/data/sampler.py#L18

        # NeighborSampler is used to create random bipartite graph between 
        # a given node and its neighbors using random walk.
        # Those random subgraphs will be used to train the graph convolution model.

        return NeighborSampler(
           self.data.edge_index,
           # the nodes that should be considered for sampling.
           node_idx=data[f'{stage}_mask'],
           # -1 indicates all neighbors will be selected
           sizes=self._num_layers * [-1],
        )

    def train_dataloader(self):
        return self.create_neighbor_sampler(stage="train")

    def validation_dataloader(self):
        return self.create_neighbor_sampler(stage="val")

    def test_dataloader(self):
        return self.create_neighbor_sampler(stage="test")