Embedding extraction from HeteroGraph Link prediction

[freh-g]

Hi everyone,

I am a PhD student I am trying to build a model for link prediction that optimizes the embedding for that specific task.
Ideally I'd like to train the model for a link prediction task and at every iteration I'd like to optimize the node features and be able at the end of the training to extract them. I am following the example available in the repository with the Hetero Graphs, the Idea is to extract the node representation from SAGEconv layer when the training is done, but I have no idea on how to do it, thank you very much for any reply.

Cheers

[rusty1s]

For link prediction tasks, you usually have an encoder and a decoder. The encoder is responsible for creating node embeddings. As such, if you are only interested in obtaining the node embeddings, you can simply do encoder(data.x, data.edge_index) (or alternatively via NeighborLoader in the mini-batch case.

[freh-g]

Thank you very much for the fast reply! So I ended up adding the following module (the part between the comments):

class Model(torch.nn.Module):
    def __init__(self, hidden_channels):
        super().__init__()
        self.encoder = GNNEncoder(hidden_channels, hidden_channels)
        self.encoder = to_hetero(self.encoder, data.metadata(), aggr='mean')
        self.decoder = EdgeDecoder(hidden_channels)

    def forward(self, x_dict, edge_index_dict, edge_label_index):
        z_dict = self.encoder(x_dict, edge_index_dict) 
        return self.decoder(z_dict, edge_label_index)

    **#### Start of the added module
    def return_embeddings(self, x_dict, edge_index_dict, edge_label_index):
        z_dict = self.encoder(x_dict, edge_index_dict)
        return z_dict
    #### End of the added module**

the GNN encoder is as github code:

class GNNEncoder(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = SAGEConv((-1, -1), hidden_channels) # Aggregation of neighbors features 
        self.conv2 = SAGEConv((-1, -1), out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index)
        return x

It takes as input not the data.x but a dictionary (x_dict) that is obtained as train_data.x_dict
Train_data is obtained through the RandomLinkSplit function

train_data, val_data, test_data = torch_geometric.transforms.RandomLinkSplit(
    num_val=0.1,
    num_test=0.1,
    neg_sampling_ratio=0.0,
    edge_types=[('protein', 'association', 'disease')],
    rev_edge_types=[('disease', 'rev_association', 'protein')],
)(data)

So even feeding the encoder with(self, x_dict, edge_index_dict, edge_label_inde) should work isn't it? Or I am missing something? I am obtaining the embeddings in the train function at every iteration as follows.

def train():
    model.train() # sets the model in training mode 
    optimizer.zero_grad()
    pred = model(train_data.x_dict, train_data.edge_index_dict,
                 train_data['protein', 'disease'].edge_label_index)#The notation model(something) is calling the forward function 
    # Add here The computations for the embeddings

    **embdict = model.return_embeddings(train_data.x_dict, train_data.edge_index_dict,
                 train_data['protein', 'disease'].edge_label_index)
    prot=embdict['protein']
    dis=embdict['disease']**
    
    # End here the computations for the embeddings

    target = train_data['protein', 'disease'].edge_label
    loss = weighted_mse_loss(pred, target)
    loss.backward()
    optimizer.step()
    return (float(loss),embdict)

Thank you Matthias!

[rusty1s]

Yes, exactly. That's what I was describing :)

[freh-g]

Thanks, you were very clear!

[karan96]

Greetings @rusty1s ,
I was following your conversation above. Given that I want to extract node embeddings of two different types of nodes in a hetero data, there are few things that I couldn't implement from above and had confusion. The dataset is as follows:

HeteroData(
  location={ node_id=[71] },
  experts={ node_id=[13631] },
  skills={ node_id=[69679] },
  (location, of, experts)={ edge_index=[2, 13631] },
  (skills, of, experts)={ edge_index=[2, 135690] },
  (experts, rev_of, location)={ edge_index=[2, 13631] },
  (experts, rev_of, skills)={ edge_index=[2, 135690] }
)

and after splitting the dataset into the train, valid and test datasets, the train looks like: -

HeteroData(
  location={ node_id=[71] },
  experts={ node_id=[13631] },
  skills={ node_id=[69679] },
  (location, of, experts)={
    edge_index=[2, 8111],
    edge_label=[3476],
    edge_label_index=[2, 3476]
  },
  (skills, of, experts)={ edge_index=[2, 135690] },
  (experts, rev_of, location)={ edge_index=[2, 8111] },
  (experts, rev_of, skills)={ edge_index=[2, 135690] }
)

I want the graph node embeddings that should look like: location [71, embedding_size] and skills [13631, embedding_size]
The parameter that is passed in the above code:- train_data.x_dict results in None for me. Is there a part that I am missing while running the above code? Could you please direct me towards any tutorial or suggest how to proceed for this?
Thanks,
Karan

[rusty1s]

Inside your model, you would have something like

self.emb_dict = torch.nn.ModuleDict({
   key: Embedding(node_id.numel(), num_features)
   for key, node_id in data.node_id_dict.items() 
})

and then call it via

x_dict = {key: emb.weight for key, emb in self.emb_dict.items()}
hetero_model(x_dict, edge_index_dict)