How to use Dataloader batches in training loop for GNN?

[adnan0819]

Hi, I am new to PyG and I can train simple GNN and GCN models. I am seeking some guidance. But for my task right now, I am doing the following:

1. Create a graph from binary image (I have a dataset of 100s of such images). Let's assume 500.
2. Therefore, assume I create 500 graphs (each has around 100-150 nodes).
3. Now, I create the data in the following manner

...
data_list = [graph1, graph2, ....]
loader = DataLoader(data_list, batch_size=2)
batch = next(iter(loader))
...

So essentially if I print out batch, I get something as follows:

DataBatch(x=[134], edge_index=[2, 1432], coord=[2], areas=[134], batch=[134], ptr=[3])
DataBatch(x=[134], edge_index=[2, 1432], coord=[2], areas=[134], batch=[134], ptr=[3])
...
DataBatch(x=[134], edge_index=[2, 1432], coord=[2], areas=[134], batch=[134], ptr=[3])

4. Suppose I am using the following code from the official tutorial to make a simple GCN.

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv

class GCN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = GCNConv(dataset.num_node_features, 16)
        self.conv2 = GCNConv(16, dataset.num_classes)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index

        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index)

        return F.log_softmax(x, dim=1)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = GCN().to(device)
data = dataset[0].to(device) 
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

model.train()
for epoch in range(200):
    optimizer.zero_grad()
    out = model(data)
    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()

@rusty1s My question is: how can I use those Databatches I showed in point 3 for training? In many examples like here, only one graph of a full dataset is used like this data = dataset[0] in Cora dataset (and also in the example above). But this is using only one graph of the full dataset. How can I use the full Dataloader's Data objects (all the graphs) in the training process?

Thank you very much for your time and will really appreciate some guidance.

PS. I tried making one gigantic graph but for my purpose "node 1", "node 2" etc. in graph 1 has a certain semantic meaning which means the same in another graph, say graph 2s "node 1", "node 2" etc. So, if I have 20 nodes in graph 1, then having nodes 21, 22, 23 etc. in the same graph changes the meaning of the problem.

[EdisonLeeeee]

Hi @adnan0819 May I ask what your downstream task is? It seems that you meant to do a graph classification/regression task.

If I understood it correctly, you can take a look at this example. More specifically, you can perform message passing on a batched graph:

for batch in loader:
    out = model(batch.x, batch.edge_index, batch.batch)

[adnan0819]

Thank you very much @EdisonLeeeee . I will try it out and report back. The downstream task is edge labeling/classification. The nodes represent buildings, the edges have attributes/weights ( representing distance) and the edge label represents if two nodes should have some similar labeling or not (if edge label between node $u$ and $v$ is $0$ it means they will not be related, if $1$ then they are related - this is what I am predicting). So the edges via topology has some attributes/weights (=distances) and a label on top of that for training. If you have any further suggestion, I will be truly grateful. Thanks again.

[EdisonLeeeee]

Sorry, but I didn't find an attribute indicating the edge label/weight in a batch in your provided code snippet:

DataBatch(x=[134], edge_index=[2, 1432], coord=[2], areas=[134], batch=[134], ptr=[3])

BTW, an edge-level task can be solved well by Graph Autoencoder (GAE). All you need to do is:

for batch in loader:
    embedding = encoder(batch.x, batch.edge_index)
    predict_label = decoder(embedding, ...)

where an encoder can be any GNN (e.g., GCN) and a decoder can be an MLP network. Hope it helps you :)