Customize collate function for torch_geometric.loader.DataLoader

[fedeotto]

In the perspective of contrastive learning applications, I would like my DataLoader to return a tuple of Batch objects according to some positives/negatives mining technique. By using standard torch.utils.DataLoader this can be accomplished by a custom collate_fn like:

def collate(batch):
    
    pos_list, neg_list = get_pairs(batch)

    return (Batch.from_data_list(batch), 
            Batch.from_data_list(pos_list), 
            Batch.from_data_list(neg_list))

Is there anything I can do to retrieve a similar thing by using PyG?

[rusty1s]

I think you can do the same for PyG. Use a torch.utils.DataLoader with custom collate_fn.

[songsong0425]

Sorry for updating this thread. I have a question about the DataLoader with the custom dataset.
When I tried to run the dataloader using from torch_geometric.loader import DataLoader, it couldn't call the batch of the graph.
For example, when I created a random graph and dataloader as below:

# Example graph data generation
bagraph = barabasi_albert_graph(num_nodes=100, num_edges=80)
example_feature = torch.rand(100, 256)
example_graph = Data(edge_index=bagraph, x=example_feature)
example_graph
# Data(x=[100, 256], edge_index=[2, 1394])  <- why it returns 1394 edges while I set num_edges=80?

# train/val/test separation
batch_size = 32
transform = RandomLinkSplit(num_val=0.1, num_test=0.2, is_undirected=True, split_labels=True)
train_data, val_data, test_data = transform(example_graph)
train_data
# Data(x=[100, 256], edge_index=[2, 978], pos_edge_label=[489], pos_edge_label_index=[2, 489], neg_edge_label=[489], neg_edge_label_index=[2, 489])

train_loader = DataLoader(train_data, shuffle=True, batch_size=batch_size)
val_loader = DataLoader(val_data, shuffle=True, batch_size=batch_size)
test_loader = DataLoader(test_data, shuffle=False, batch_size=batch_size)

the data loaders couldn't get mini-batches with KeyError.

for batch in loader:
    # Access the batch data
    batch_x = batch.x
    batch_edge_index = batch.edge_index

# KeyError: 2

Full error message

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Cell In [17], line 1
----> 1 for batch in train_loader:
      2     # Access the batch data
      3     batch_x = batch.x
      4     batch_edge_index = batch.edge_index

File ~/.venv/torch1.12.1-py3.8-cuda11.3/lib/python3.8/site-packages/torch/utils/data/dataloader.py:681, in _BaseDataLoaderIter.__next__(self)
    678 if self._sampler_iter is None:
    679     # TODO(https://github.com/pytorch/pytorch/issues/76750)
    680     self._reset()  # type: ignore[call-arg]
--> 681 data = self._next_data()
    682 self._num_yielded += 1
    683 if self._dataset_kind == _DatasetKind.Iterable and \
    684         self._IterableDataset_len_called is not None and \
    685         self._num_yielded > self._IterableDataset_len_called:

File ~/.venv/torch1.12.1-py3.8-cuda11.3/lib/python3.8/site-packages/torch/utils/data/dataloader.py:721, in _SingleProcessDataLoaderIter._next_data(self)
    719 def _next_data(self):
    720     index = self._next_index()  # may raise StopIteration
--> 721     data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
    722     if self._pin_memory:
    723         data = _utils.pin_memory.pin_memory(data, self._pin_memory_device)

File ~/.venv/torch1.12.1-py3.8-cuda11.3/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py:49, in _MapDatasetFetcher.fetch(self, possibly_batched_index)
     47 def fetch(self, possibly_batched_index):
     48     if self.auto_collation:
---> 49         data = [self.dataset[idx] for idx in possibly_batched_index]
     50     else:
     51         data = self.dataset[possibly_batched_index]

File ~/.venv/torch1.12.1-py3.8-cuda11.3/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py:49, in <listcomp>(.0)
     47 def fetch(self, possibly_batched_index):
     48     if self.auto_collation:
---> 49         data = [self.dataset[idx] for idx in possibly_batched_index]
     50     else:
     51         data = self.dataset[possibly_batched_index]

File ~/.venv/torch1.12.1-py3.8-cuda11.3/lib/python3.8/site-packages/torch_geometric/data/data.py:457, in Data.__getitem__(self, key)
    456 def __getitem__(self, key: str) -> Any:
--> 457     return self._store[key]

File ~/.venv/torch1.12.1-py3.8-cuda11.3/lib/python3.8/site-packages/torch_geometric/data/storage.py:104, in BaseStorage.__getitem__(self, key)
    103 def __getitem__(self, key: str) -> Any:
--> 104     return self._mapping[key]

KeyError: 2

It seems like that loader can't separate the graph and I guess it needs collate_fn.
But doesn't DataLoader automatically call the innate collate_fn? Should I set the collate function manually?

[rusty1s]

DataLoader takes in a list of data objects, while you are only inputting a single Data object. If you want to perform subgraph sampling, take a look at LinkNeighborLoader.

[songsong0425]

Thank you for your reply. Then, I'll check the LinkNeightborLoader soon. Thanks a lot!