Link Prediction on Heterogenous Graph

[anishparajuli555]

Is there any example of link prediction usage on Heterogenous Graph? I am looking forward to implement link prediction on similar task as this one: https://github.com/pyg-team/pytorch_geometric/blob/master/examples/hetero/load_csv.py

[rusty1s]

We only have an open PR yet, see #3350.

[anishparajuli555]

Thanks @rusty1s !

I have been getting this error while trying to implement the example. Do you have any idea what might have caused this issue?

[rusty1s]

Can you show me the full model?

[anishparajuli555]

from torch_geometric.nn import SAGEConv, to_hetero

class EdgeClassifier(torch.nn.Module):
    def __init__(self, hidden_channels, edge_labels):
        super().__init__()
        # Set up GNN for message passing
        self.encoder = GNN(hidden_channels=hidden_channels, out_channels=hidden_channels)
        # Make the GNN model heterogeneous
        self.encoder = to_hetero(self.encoder, data.metadata(), aggr='sum')

        # DistMult based "decoder" for edge classification
        self.decoder = EdgeDecoder(hidden_channels, edge_labels)

    def forward(self, x_dict, edge_index_dict, edge_label_index):
        # Perform message passing
        z = self.encoder(x_dict, edge_index_dict)

        # The edges that go into the decoder are the supervision edges (edge_label_index)
        out = self.decoder(z, edge_label_index)
        return(out)

from torch.nn import Parameter

class EdgeDecoder(torch.nn.Module):
    def __init__(self, hidden_channels, edge_labels):
        super().__init__()

        self.edge_labels = edge_labels
        num_relations = torch.numel(self.edge_labels)

        self.rel_emb = Parameter(torch.Tensor(num_relations, hidden_channels))
        self.reset_parameters()

    def reset_parameters(self):
        torch.nn.init.xavier_uniform_(self.rel_emb)


    def decode(self, z, edge_index, edge_type):
        # We only need to decode on the ('user','rates','movie') edges
        z_src = z['user']
        z_dst = z['movie']
        z_src = z_src[edge_index[0]]
        z_dst = z_dst[edge_index[1]]
        rel = self.rel_emb[edge_type]
        return torch.sum(z_src * rel * z_dst, dim=1)

    def forward(self, z, edge_label_index):
        # Initialise scores array, which records the DistMult score for
        # each supervision edge and edge_label (rating)
        scores = torch.zeros(edge_label_index.shape[1], self.edge_labels.shape[0])
        for label in self.edge_labels:
            score = self.decode(z,edge_label_index, label)
            scores[:,label] = score
        # Return log probabilities for each class
        return(F.log_softmax(scores, dim=0))

class GNN(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = SAGEConv((-1, -1), hidden_channels)
        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

[rusty1s]

Thanks a lot. I cannot find any error within your model, so I expect the error to be in the input. Can you confirm that edge_index_dict only contains LongTensors?

[gulycat1214]

This is likely because x is expected to be a torch.FloatTensor and somehow you are passing a torch.LongTensor