What're the standard ways of padding data for GNNs?

Question

I am working on materials property prediction using GNNs with torch_geometric.

Each data in my dataset has different number of feature vectors x, edge_index vectors edge_index. What I'm doing is to pad all the feature vectors and edge_index features, respectively, to the length of the element with maximum length. For example, this is what I do for the edge_index vector:

In:  edge_index = train_set[14].edge_index
     edge_index
Out: tensor([[1, 0, 2, 0, 2, 1],
        [0, 1, 0, 2, 1, 2]])
In:  fill_zeros(edge_index, 15, 1)
Out: tensor([[1, 0, 2, 0, 2, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 1, 0, 2, 1, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0]])

i.e., I defined a function called fill_zeros to fill in the dummy variables 0's till the tensor being acted on has length $15$.

However, this incurs many problems when feeding my data into GNN models when doing the message passing. More specifically, when running propagate(). I wonder what're the standard ways of padding data for GNN models to avoid such problems?

Answer 1

One way is to avoid padding by batching.

The idea is that you process data in batches where each Data in the batch has the same amount of feature vectors.

In pytorch_geometric from torch_geometric.data import Data, DataLoader

Step 1

Prepare your graph dataset as a list of Data objects. Each Data object should contain a consistent number of nodes and edges.

dataset = [Data, Data, ...]

Step 2

Sort the Data objects in your dataset based on the number of features you want to use for batching. In your case you mentioned edges, so you can use sort like this

dataset.sort(key=lambda data: data.num_edges)

Step 3 Create batches where each batch contains data with the same number of features. You can do this by iterating through the sorted dataset and grouping together graphs with the same feature size.

batch_size = 32  # Adjust the batch size as needed
batched_data = []
current_batch = []

for data in dataset:
    # Check if the current data has the same feature size as the data in the current batch
    if not current_batch or data.num_nodes == current_batch[0].num_nodes:
        current_batch.append(data)
    else:
        batched_data.append(current_batch)
        current_batch = [data]

# Append the last batch
if current_batch:
    batched_data.append(current_batch)

Step 4

Create a DataLoader from the batched_data list

data_loader = DataLoader(batched_data, batch_size=None, shuffle=True)

note: batch_size is set to None because you've already grouped the data into batches.

With this setup, each batch in the DataLoader will contain Data objects with the same number of features, ensuring that you don't need to perform any padding.