Using edge attributes with HeteroData and splitting the dataset

[adirsolo]

Hi all,

I am new to pytorch_geometric and I am encountering some issues that I need assistance with:

1. I would like to approach my problem as a link prediction task, where a user's purchase of an item (referred to as a "deal" in my case) is involved. To accomplish this, I have a binary label column and I intend to consider the values of 1 as positive edges and the values of zero as negative edges.
2. My objective is to utilize different attributes as edge attributes, such as 'money', 'deal_discount', and 'deal_commission', employing the GATConv module.
3. I need to divide my dataset based on the integer values present in a column named "month" instead of employing RandomLinkSplit. Specifically, I want to allocate the months 1-4 for training purposes and designate the fifth month as the test set.

I have attached a sample notebook for reference. I would greatly appreciate any assistance, including code examples.

Thank you!
GNN_Toy_Example.zip

[rusty1s]

Hey, I haven't looked at your notebook yet, but this all sounds correct to me. Do you have any specific questions you would like to have an answer for?

[adirsolo]

I am encountering multiple errors when trying to perform the abovementioned tasks.
Can you provide code snippets/references to assist me with each task?
Thanks!

[rusty1s]

Sorry for late reply. Can you give some more details about the errors? This would help me to answer your questions :)

[adirsolo]

I want to split my dataframe to train and test graphs based on a column values (in my case, the month column).
I encountered the following error:

IndexError Traceback (most recent call last)
/usr/local/lib/python3.10/dist-packages/torch_geometric/nn/conv/message_passing.py in _lift(self, src, edge_index, dim)
265 index = edge_index[dim]
--> 266 return src.index_select(self.node_dim, index)
267 except (IndexError, RuntimeError) as e:

IndexError: index out of range in self

During handling of the above exception, another exception occurred:

IndexError Traceback (most recent call last)
7 frames
/usr/local/lib/python3.10/dist-packages/torch_geometric/nn/conv/message_passing.py in _lift(self, src, edge_index, dim)
267 except (IndexError, RuntimeError) as e:
268 if index.min() < 0 or index.max() >= src.size(self.node_dim):
--> 269 raise IndexError(
270 f"Encountered an index error. Please ensure that all "
271 f"indices in 'edge_index' point to valid indices in "

IndexError: Encountered an index error. Please ensure that all indices in 'edge_index' point to valid indices in the interval [0, 1] (got interval [0, 2])

My code:

def create_graph_data(df):
  node_labels = df['Source'].append(df['Target']).unique()
  node_indices = {label: index for index, label in enumerate(node_labels)}
  x = torch.tensor(df[['NodeFeature1', 'NodeFeature2']].values, dtype=torch.float)
  y = torch.tensor(pd.DataFrame([df['Label_1'].values, df['Label_2'].values]).T.values, dtype=torch.float)

  edge_index = []
  for source, target in zip(df['Source'], df['Target']):
      source_idx = node_indices[source]
      target_idx = node_indices[target]
      edge_index.append([source_idx, target_idx])
  edge_index = torch.tensor(edge_index, dtype=torch.long).t().contiguous()

  edge_attr = torch.tensor(df[['EdgeFeature1', 'EdgeFeature2']].values, dtype=torch.float)
  data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr, y=y)
  return data

import torch
import pandas as pd
from torch_geometric.data import Data
from torch_geometric.nn import GCNConv, GATConv

df = pd.DataFrame({'Source': ['A', 'B', 'C', 'A', 'D'],
                   'Target': ['B', 'C', 'A', 'C', 'A'],
                   'EdgeFeature1': [0.1, 0.2, 0.3, 0.9, 0.6],
                   'EdgeFeature2': [0.6, 0.7, 0.8, 0.5, 0.4],
                   'NodeFeature1': [0.2, 0.4, 0.6, 0.1, 0.2],
                   'NodeFeature2': [0.2, 0.4, 0.6, 0.8, 0.3],
                   'Month':[4,4,4,5,5],
                   'Label_1': [0, 1, 0, 1, 1],
                   'Label_2': [1, 1, 0, 0, 1]})

df_train = df[df['Month']<5]
df_test = df[df['Month']==5]

data_train = create_graph_data(df_train)
data_test = create_graph_data(df_test)

class GCN(torch.nn.Module):
    def __init__(self, in_channels, edge_attr_dim, hidden_channels):
        super(GCN, self).__init__()
        self.conv1 = GATConv(in_channels=in_channels, out_channels=hidden_channels)
        self.conv2 = GATConv(in_channels=hidden_channels, out_channels=2)

    def forward(self, data):
        x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr
        # x = torch.cat([x, edge_attr], dim=1)
        x = self.conv1(x, edge_index)
        x = torch.relu(x)
        x = self.conv2(x, edge_index)
        x = torch.sigmoid(x)
        return x

model = GCN(in_channels=2, edge_attr_dim=2, hidden_channels=16)

optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
criterion = torch.nn.BCEWithLogitsLoss()

for epoch in range(100):
    model.train()
    optimizer.zero_grad()
    out = model(data_train)
    loss = criterion(out.squeeze(), data_train.y)
    loss.backward()
    optimizer.step()

model.eval()
with torch.no_grad():
    pred = model(data_test).squeeze()

print(pred)

[adirsolo]

My main question is how to adapt RandomLinkSplit to work not randomly but based on a custom condition for the Transductive link prediction task. For instance, month 4 will be used as a train set and month 5 as a test set.

[rusty1s]

In that case, you can do something like

test_mask = data.month == 5
train_mask = ~test_mask

If month is a node-level attribute, you would need to lift it to an edge-level attribute afterwards via

test_mask = test_mask[edge_index[0]]
train_mask = train_mask[edge_index[0]]