Edge attributes with heterographs

[GianlucaDeStefano]

Hi all, I am trying to train a model that performs link prediction on an Heterograph. In particular, at the moment I am struggling with associating to every edge and additional embedding.

Currently, I am doing this:

self.edge_type_emb = torch.nn.Embedding(32, 128)

self.conv = GATv2Conv(128, 32,4, add_self_loops=False, edge_dim=128)
self.conv = to_hetero(self.gnn, metadata=dataset.data.metadata(), aggr='sum')


ft_edges = {}

for edge_name, edge_types in batch.edge_types_dict.items():
    edge_types_embedding = self.edge_type_emb(edge_types)
    ft_edges[edge_name] = edge_types_embedding

self.conv(node_features_dict, batch.edge_index_dict, ft_edges)

Where node_features_dict and edge_attr_dict are two vocabularies containing for every node and edge type a feature vector of shape [n, 128].
However, by doing this I get a dimensionality exception in the internal computations of the convolution:

Cell In[85], line 276, in GATv2Conv.message(self, x_j, x_i, edge_attr, index, ptr, size_i)
    272     print(x.shape)
    273     print(edge_attr.shape)
--> 276     x = x + edge_attr
    278 x = F.leaky_relu(x, self.negative_slope)
    279 alpha = (x * self.att).sum(dim=-1)

RuntimeError: The size of tensor a (3284) must match the size of tensor b (1466) at non-singleton dimension 0`

Essentially there the dimension of x and edge_attr are respectively:
x: torch.Size([3284, 4, 32]) edge_attr: torch.Size([1466, 4, 32])

Am I doing something wrong? Is there an example showing how to handle edge-related attributes with Heterographs?

[wsad1]

self.conv = GATv2Conv(128, 32,4, add_self_loops=False, edge_dim=128)
self.conv(node_features_dict, batch.edge_index_dict, edge_attr_dict)

I am a little confused did you use to_hetero to convert GATv2Conv to its heterogenous equivalent? Because the code above won't work as GATv2Conv expects Tensor and not dictionaries like node_features_dict.

[GianlucaDeStefano]

Yes, sorry I cleaned a bit the code and forgot to add that line back

[GianlucaDeStefano]

I checked and seems to be a problem of the batch generation.
Essentially my dataset looks like this:

HeteroData(
  SYNTAX_NODE={
    node_id=[29895],
    type_id=[29895]
  },
  SYNTAX_TOKEN={
    node_id=[26626],
    value_id=[26626],
    type_id=[26626]
  },
  (SYNTAX_NODE, edge_type_1, SYNTAX_NODE)={
    edge_index=[2, 29893],
    edge_attr=[29893]
  },
  (SYNTAX_NODE, edge_type_1, SYNTAX_TOKEN)={
    edge_index=[2, 26626],
    edge_attr=[26626]
  },
  (SYNTAX_NODE, edge_type_2, SYNTAX_NODE)={
    edge_index=[2, 5989],
    edge_attr=[5989]
  },
  (SYNTAX_NODE, edge_type_4, SYNTAX_NODE)={
    edge_index=[2, 5533],
    edge_attr=[5533]
  },
  (SYNTAX_NODE, edge_type_2, SYNTAX_TOKEN)={
    edge_index=[2, 120],
    edge_attr=[120]
  },
  (SYNTAX_TOKEN, edge_type_2, SYNTAX_NODE)={
    edge_index=[2, 182],
    edge_attr=[182]
  },
  (SYNTAX_TOKEN, edge_type_2, SYNTAX_TOKEN)={
    edge_index=[2, 12],
    edge_attr=[12]
  },
  (SYNTAX_NODE, edge_type_3, SYNTAX_NODE)={
    edge_index=[2, 1121],
    edge_attr=[1121]
  },
  (SYNTAX_NODE, edge_type_3, SYNTAX_TOKEN)={
    edge_index=[2, 1],
    edge_attr=[1]
  }
...

So inside the dataset, the number of edge_indexes and number o edge_attributes is the same.

However, once I sample a batch there are some inconsistencies between the number of edges sampled and their relative attributes and this is the cause of the problem:

HeteroData(
  SYNTAX_NODE={
    node_id=[1805],
    type_id=[1805]
  },
  SYNTAX_TOKEN={
    node_id=[2],
    value_id=[2],
    type_id=[2]
  },
  (SYNTAX_NODE, edge_type_1, SYNTAX_NODE)={
    edge_index=[2, 3245],
    edge_types=[1440]
  },
  (SYNTAX_NODE, edge_type_1, SYNTAX_TOKEN)={
    edge_index=[2, 2],
    edge_types=[2]
  },
  (SYNTAX_NODE, edge_type_2, SYNTAX_NODE)={
    edge_index=[2, 2148],
    edge_types=[343]
  },
  (SYNTAX_NODE, edge_type_4, SYNTAX_NODE)={
    edge_index=[2, 2018],
    edge_types=[213]
  },
  (SYNTAX_NODE, edge_type_3, SYNTAX_NODE)={
    edge_index=[2, 1853],
    edge_types=[48],
    input_id=[64],
    edge_label_index=[2, 128],
    edge_label=[128]
  },
  (SYNTAX_NODE, edge_type_2, SYNTAX_TOKEN)={
    edge_index=[2, 2],
    edge_types=[2]
  },
  (SYNTAX_TOKEN, edge_type_2, SYNTAX_NODE)={
    edge_index=[2, 2],
    edge_types=[2]
  },
  (SYNTAX_TOKEN, edge_type_2, SYNTAX_TOKEN)={
    edge_index=[2, 2],
    edge_types=[0]
  },
  (SYNTAX_NODE, edge_type_3, SYNTAX_TOKEN)={
    edge_index=[2, 0],
    edge_types=[0]
  }
...

I am using the LinkNeighborLoader:

train_loader = LinkNeighborLoader(
    data= training_data,
    num_neighbors=[nodes_to_sample_per_layer]*num_message_passing_operations,
    neg_sampling_ratio=neg_edges_sampling_ratio,
    edge_label_index=(("SYNTAX_NODE", "edge_type_3", "SYNTAX_NODE"), edge_index_training),
    edge_label=edge_label,
    batch_size=batch_size,
    shuffle=True,
    drop_last = True
)

But I really can't understand why this happens.

[wsad1]

Also in the example you shared

  (SYNTAX_NODE, edge_type_1, SYNTAX_TOKEN)={
    edge_index=[2, 2],
    edge_types=[2]
  },

How did the edge_types variable get into the batch?
Also it tried out the below code to see if edge_attr looks ok in the batch object and it works fine

import torch
from torch_geometric.data import HeteroData

def get_edge_index(num_src_nodes, num_dst_nodes, num_edges):
    row = torch.randint(num_src_nodes, (num_edges, ), dtype=torch.long)
    col = torch.randint(num_dst_nodes, (num_edges, ), dtype=torch.long)
    return torch.stack([row, col], dim=0)

def hetero_data():
    data = HeteroData()
    data['paper'].node_id = torch.randn(8, 16)
    data['author'].node_id = torch.randn(10, 8)

    data['paper', 'paper'].edge_index = get_edge_index(8, 8, 10)
    data['paper', 'paper'].edge_attr = torch.randn(10, 16)
    data['author', 'paper'].edge_index = get_edge_index(10, 8, 10)
    data['paper', 'author'].edge_index = get_edge_index(8, 10, 10)
    data['paper', 'author'].edge_attr = torch.randn(10, 8)

    return data

from torch_geometric.loader import LinkNeighborLoader
loader = LinkNeighborLoader(hetero_data(), num_neighbors=[5, 5],
                edge_label_index=(('paper', 'to', 'author'), None),
                shuffle=True, num_workers=0)
print(next(iter(loader)))

This prints

HeteroData(
  paper={ node_id=[5, 16] },
  author={ node_id=[4, 8] },
  (paper, to, paper)={
    edge_index=[2, 5],
    edge_attr=[5, 16]
  },
  (author, to, paper)={ edge_index=[2, 6] },
  (paper, to, author)={
    edge_index=[2, 3],
    edge_attr=[3, 8],
    input_id=[1],
    edge_label_index=[2, 1]
  }
)

[GianlucaDeStefano]

I create the edge_type in the following manner:

edge_types = torch.tensor([self.get_edge_type_id(e.type) for e in edges])
heteroData[node_type_1, edge_name, node_type_2].edge_types = edge_types

So edge_types is just a tensor of integers that then are converted to a vectorial representation using an Embedding layer (I updated the code snipped above).

I then feed to the GNN a dictionary containing the association 'edge_name' -> 'edge_types_embedding'

[wsad1]

Could you share the the shapes of the tensors in the HeteroData object i.e the output of print(your_data). Because as I showed in the example if the edge_attr is of the right shape in the original data it should look ok in the batch.

[GianlucaDeStefano]

Sure here is the raw print of my HeteroData Object that I give to the Dataloader:

HeteroData(
  SYNTAX_NODE={
    node_id=[267657],
    type_id=[267657]
  },
  SYNTAX_TOKEN={
    node_id=[254028],
    value_id=[254028],
    type_id=[254028]
  },
  (SYNTAX_NODE, edge_type_1, SYNTAX_NODE)={
    edge_index=[2, 267590],
    edge_types=[267590]
  },
  (SYNTAX_NODE, edge_type_1, SYNTAX_TOKEN)={
    edge_index=[2, 254028],
    edge_types=[254028]
  },
  (SYNTAX_NODE, edge_type_2, SYNTAX_NODE)={
    edge_index=[2, 36284],
    edge_types=[36284]
  },
  (SYNTAX_NODE, edge_type_3, SYNTAX_NODE)={
    edge_index=[2, 29424],
    edge_types=[29424]
  },
  (SYNTAX_NODE, edge_type_4, SYNTAX_NODE)={
    edge_index=[2, 1594],
    edge_types=[1594]
  },
  (SYNTAX_NODE, edge_type_5, SYNTAX_NODE)={
    edge_index=[2, 4],
    edge_types=[4]
  },
  (SYNTAX_NODE, edge_type_6, SYNTAX_NODE)={
    edge_index=[2, 5],
    edge_types=[5]
  },
  (SYNTAX_NODE, edge_type_7, SYNTAX_NODE)={
    edge_index=[2, 7192],
    edge_types=[7192]
  },
  (SYNTAX_NODE, edge_type_2, SYNTAX_TOKEN)={
    edge_index=[2, 203],
    edge_types=[203]
  },
  (SYNTAX_TOKEN, edge_type_2, SYNTAX_NODE)={
    edge_index=[2, 360],
    edge_types=[360]
  },
    (SYNTAX_TOKEN, edge_type_2, SYNTAX_TOKEN)={
    edge_index=[2, 8],
    edge_types=[8]
  },
  (SYNTAX_NODE, edge_type_7, SYNTAX_TOKEN)={
    edge_index=[2, 13],
    edge_types=[13]
  }
)

I am currently debugging, I will let you know if I find the reason

[GianlucaDeStefano]

Ok I found what the problem was, among the transformations I was using there was T.AddSelfLoops() that was adding edges without attributes :(. Removing that everything works fine. Really an obvious and stupid mistake