Add sampler argument to NeighborLoader.

[wsad1]

🚀 Feature

Add a sampler argument to NeighborLoader. The sampler would return a list of nodes to load.

Motivation

Tow keys motivations.

1. Make distributed_sampling look much more similar to distributed_batching.
2. I think this is need for multi gpu training with pytorch lightning, correct me if I am wrong but the example here won't split the data across multiple gpus as a distributed sampler is needed for that.

[rusty1s]

Thanks in advance for the discussion. Do you have a small example to showcase what you mean exactly?

1. Aren't both interfaces already quite similar. In distributed_batching, we first need to split the dataset according to process, while in distributed_sampling, we split the input nodes according to process.
2. The PL example will split the data already AFAIK. This is done internally by PL, which accesses the loader.dataset property, splits them across processes, and re-instantiates the loaders.

[wsad1]

Just realized that NeighborLoader inherits torch.utils.data.DataLoader, so you can pass a sampler into it if needed. So, this feature isn't needed. But I do have some other thoughts on the above two points.

1. I agree that they are quite similar. But do you think using NeighborLoader with DistributedSampler would make it look exactly the same, plus we could avoid using NeighborSampler.
2. You are right it does split up the data, internally using DistributedSampler as described here. Here also do you think moving to NeighborLoader is a good idea?

[rusty1s]

Are you referring to the point that NeighborSampler requires gathering (sub-)feature matrices in a post-processing step? I think we want to replace NeighborSampler with NeighborLoader in the long run, but for distributed training, it currently does not matter as all graph data needs to exist on every machine anyway. Something we need to tackle in a new release.

[wsad1]

@rusty1s , sorry for the later reply.

Yes I was referring to that point and the fact that NeighborSampler can't use a DistributedSampler and instead the data is split using this line.

Why do you say all graph data needs to exist on every machine? Isn't the data moved to the gpu when this line is executed? Let me know if I am missing something.

[rusty1s]

Yeah, i was referring to distributing across multiple nodes. Here, each node/machine currently still needs the full data.

[schlyah]

Is it possible to pass a DistributedSampler to NeighborLoader as it is the case for torch.utils.data.DataLoader ? When I tried to do that the batches only contain a few elements of the original data.

[schlyah]

Thanks for your reply.

The following code:

    dataset = PygNodePropPredDataset(name='ogbn-arxiv')
    data = dataset[0]
    data.num_classes = dataset.num_classes
    data.relations = torch.full(size=(data.num_edges, 1), fill_value=args.binary_preactivation)
    split_idx = dataset.get_idx_split()
    train_idx = split_idx['train']
    batch_size = 100000

    # Using distributed sampler
    train_sampler = torch.utils.data.distributed.DistributedSampler(data, num_replicas=2, rank=rank)
    train_loader1 = NeighborLoader(data,
                                   num_neighbors=[args.sampling_neighbor_size] * args.num_layers_sampling,
                                   input_nodes=train_idx,
                                   batch_size=batch_size,
                                   num_workers=args.num_workers,
                                   sampler=train_sampler)

    # Using index
    idx = train_idx.split(train_idx.size(0) // 2)[rank]
    train_loader2 = NeighborLoader(data,
                                   num_neighbors=[args.sampling_neighbor_size] * args.num_layers_sampling,
                                   shuffle=True,
                                   input_nodes=idx,
                                   batch_size=batch_size,
                                   num_workers=args.num_workers)
    for batch in train_loader1:
        print(("TL1", gpu, len(batch.x), batch.batch_size))

    for batch in train_loader2:
        print(("TL2", gpu, len(batch.x), batch.batch_size))

Produces this output:
('TL1', 1, 29, 4)
('TL1', 0, 16, 4)
('TL2', 0, 110580, 45470)
('TL2', 1, 110828, 45470)

[rusty1s]

I am pretty sure you would need to pass in train_idx to the DistributedSampler:

    train_sampler = torch.utils.data.distributed.DistributedSampler(train_idx.tolist(), num_replicas=2, rank=rank)

[schlyah]

Thanks for your help. That works.

[XGZepto]

I am pretty sure you would need to pass in train_idx to the DistributedSampler:

    train_sampler = torch.utils.data.distributed.DistributedSampler(train_idx.tolist(), num_replicas=2, rank=rank)

A follow up on this. How do we handle this when dealing with HeteroData? NeighborLoader takes a tuple with node_type and a list of indices but what about DistributedSampler?

[rusty1s]

Haven't tested, but I think it should work out-of-the-box as outlined here since the node_type is already specified through NeighborLoader.

[manikandanbalasubramani03]

Hi!
As you mentioned, all graph data needs to exist on every machine when training and testing the model. However, I am encountering a CUDA out of memory issue while using three GPUs, specifically with datasets like OGBN-Products.

Additionally, I have a question regarding the testing process. Did they use the entire dataset for testing? If so, won't that also lead to a CUDA out of memory issue for large datasets like OGBN-Products?

https://github.com/pyg-team/pytorch_geometric/blob/master/examples/reddit.py

[rusty1s]

Are you referring to the model.inference method? Usually, only mini-batch data sits on GPU.

[manikandanbalasubramani03]

Thanks for your reply!

Actually in the below link

https://github.com/pyg-team/pytorch_geometric/blob/master/examples/reddit.py
@torch.no_grad()
def test():
model.eval()
y_hat = model.inference(data.x, subgraph_loader).argmax(dim=-1)
y = data.y.to(y_hat.device)

accs = []
for mask in [data.train_mask, data.val_mask, data.test_mask]:
    accs.append(int((y_hat[mask] == y[mask]).sum()) / int(mask.sum()))
return accs

data.train_mask, data.val_mask, data.test_mask -> Does these values containes entire dataset ?

For eg:
RedditData(x=[232965, 602], y=[232965], adj_t=[232965, 232965, nnz=114615892], num_nodes=232965, train_mask=[232965], test_mask=[232965], val_mask=[232965])

data.train_mask ->train_mask=[232965],
data.test_mask ->test_mask=[232965],
data.val_mask ->val_mask=[232965])

Were these values used in the test process for each epoch?

If so, when it comes to data parallelism, the data is split across GPUs, and therefore, the entire dataset values are not used at once in the test process for each epoch.

If the entire dataset values are not used in the test process for each epoch, is this function working correctly? Is this approach correct?

And in my case the inference model looks like below:

def inference(self, x, device):
    x = x.to(torch.float)

    self.eval()
    x = x.to(device)
    for i, lin in enumerate(self.lins[:-1]):
        x = lin(x)
        x = self.bns[i](x)
        x = F.relu(x)
        x = F.dropout(x, p=self.dropout, training=self.training)
    x = self.lins[-1](x)
    return x

and i am not using the subgraph_loader ?
could you please explain or share me the link function of subgraph_loader?_

[rusty1s]

For data parallelism scenarios, you can and you should indeed not use this method of layer-wise inference/subgraph_loader. Instead, you can do the regular forward pass for a subset of test indices.