dataloader for dynamic graph

[Peien429]

I have some snapshots of a dynamic graph. I want to use some static gnn methods, such as GAT, to generate embeddings of nodes in each snapshot and use RNN to encode the the whole snapshots' node embeddings.

The number of nodes of the snapshots is very large and the number of snapshots is also very large. So that I can't generate all nodes' embeddings of all snapshots with GAT at the same time. So I want to generate a batch of nodes' embeddings of every snapshots with GAT and feed them into the RNN module.

How can i create batch for train or test? Specificly, how can i make sure the embeddings obtained from different snapshot belonged to the same nodes?

[rusty1s]

The best way I can think of right now is to sample a number of nodes/seed nodes, and then compute the induced subgraph for these. Alternatively, you may be able to leverage NeighborLoader for this job, which samples neighbors across snapshots for the same set of seed nodes. You can then feed this individual subgraphs into a GNN and pass them to an RNN. Hope this helps.

[rusty1s]

Yes, this is one option I thought of as well. In this setting, you would need to ensure that sampled seed nodes are the same by customizing the sampler argument of a DataLoader, see here. For example, you can create a RandomSampler based on the same underlying generator.

Another option is to wrap your loaders into another DataLoader, such that it becomes

class MyDataLoader(torch.utils.data.DataLoader):
    def __init__(self, ...):
         seed_nodes = ...
         self.loader1 = ...
         self.loader2 = ...
         super().__init__(seed_nodes, collate_fn=self.collate_fn)
         
    def collate_fn(self, seed_nodes):
         return self.loader1.collate_fn(seed_nodes), self.loader2.collate_fn(seed_nodes)

[Peien429]

Hi，I read the pytorch docs and use the following code.

import torch 

from torch.utils.data import Dataset, DataLoader

class myDataset(Dataset):
   def __init__(self):
       super().__init__()
       self.x = torch.arange(0,100).reshape(100,1)

   def __len__(self):
       return len(self.x)

   def __getitem__(self, idx):
       return self.x[idx]

if __name__ == '__main__':
   data = myDataset()
   torch.manual_seed(0)
   generator1 = torch.Generator()
   loader1 = DataLoader(data, shuffle=True, generator=generator1, batch_size=4)
   iter1 = iter(loader1)

   torch.manual_seed(0)
   generator2 = torch.Generator()
   loader2 = DataLoader(data, shuffle=True, generator=generator2, batch_size=4)
   iter2 = iter(loader2)

   for i in range(2):
       print("...")
       x = next(iter1)
       y = next(iter2)
       print(x)
       print(y)

The result is as follows:

...
tensor([[37],
       [17],
       [45],
       [24]])
tensor([[37],
       [17],
       [45],
       [24]])
...
tensor([[54],
       [19],
       [59],
       [42]])
tensor([[54],
       [19],
       [59],
       [42]])

I have known that the same seed makes generator1 and generator2 generate the same index, resulting in the loader1 and loader2 yield the same data.
I tried the following codes and it seems to be right.

   graph1 = Data(x=torch.randn(100,5), id=torch.arange(100), edge_index=torch.randint(0, 100, (2, 10000)))
   graph2 = Data(x=torch.randn(100,5), id=torch.arange(100), edge_index=torch.randint(0, 100, (2, 1000)))
   batch_size = 8

   torch.manual_seed(0)
   generator1 = torch.Generator()
   loader1 = NeighborLoader(graph1,num_neighbors=[2] * 2,batch_size=batch_size,generator=generator1,shuffle=True)
   torch.manual_seed(0)
   generator2 = torch.Generator()
   loader2 = NeighborLoader(graph2,num_neighbors=[2] * 2,batch_size=batch_size,generator=generator2,shuffle=True)
   
   iter1 = iter(loader1)
   iter2 = iter(loader2)
   for i in range(3):
       sub_graph1 = next(iter1)
       sub_graph2 = next(iter2)
       print(sub_graph1.id[:batch_size])
       print(sub_graph2.id[:batch_size])
       print("...")

The result is as follows:

tensor([37, 17, 45, 24, 54, 19, 59, 42])
tensor([37, 17, 45, 24, 54, 19, 59, 42])
...
tensor([90, 14, 77, 60, 80, 27, 78, 32])
tensor([90, 14, 77, 60, 80, 27, 78, 32])
...
tensor([18,  1, 68, 97, 55, 56,  9,  5])
tensor([18,  1, 68, 97, 55, 56,  9,  5])
...

However, I still don't know how to use "seed_nodes", could you please tell me how to use it?
Thanks a lot!

[rusty1s]

Yes, this looks perfect to me. What do you mean by using "seed nodes"?

[Peien429]

Hi! I don't know how to use your previous code. Specificly, I don't know the variable types of "seed_nodes" and "self.loader1" that I should use. Could you explain the code to me? Thanks a lot!

class MyDataLoader(torch.utils.data.DataLoader):
    def __init__(self, ...):
         seed_nodes = ...
         self.loader1 = ...
         self.loader2 = ...
         super().__init__(seed_nodes, collate_fn=self.collate_fn)
         
    def collate_fn(self, seed_nodes):
         return self.loader1.collate_fn(seed_nodes), self.loader2.collate_fn(seed_nodes)

[rusty1s]

Got it, just thought you fixed it already with the generator method. For the unified dataloader, seed_nodes refers to the nodes you want to start sampling from. In your case, this seems to be all nodes in the graph, so you can do something like seed_nodes = torch.arange(data.num_nodes):

seed_nodes = torch.arange(data.num_nodes)
self.loader1 = NeighborLoader(data, num_neighbors=...)
self.loader2 = NeighborLoader(data, num_neighbors=...)
super().__init__(seed_nodes, collate_fn=self.collate_fn)