Add code for calculation of MRR in torch_geometric/nn/kge/base.py

[songsong0425]

Hi, always thank you for your effort in updating the package.
I'd like to suggest the idea of an evaluation metric on the KGE model.

In the current version of torch_geometric/nn/kge/base.py, only two metrics for evaluation exist.
As you know, in the original papers in KGE models, such as RotatE, the mean reciprocal rank (MRR) score is calculated for the performance evaluation.
To mimic them, it would be better to add the evaluation metrics in the base.py code.
How about adding one line for this task as below?

def test(
        ...
        arange = range(head_index.numel())
        arange = tqdm(arange) if log else arange

        mean_ranks, hits_at_k = [], []
        for i in arange:
            h, r, t = head_index[i], rel_type[i], tail_index[i]

            scores = []
            tail_indices = torch.arange(self.num_nodes, device=t.device)
            for ts in tail_indices.split(batch_size):
                scores.append(self(h.expand_as(ts), r.expand_as(ts), ts))
            rank = int((torch.cat(scores).argsort(
                descending=True) == t).nonzero().view(-1))
            mean_ranks.append(rank)
            hits_at_k.append(rank < k)

        mrr = (1. / torch.tensor(mean_ranks, dtype=torch.float)).mean() # referenced from rgcn_link_pred.py file
        mean_rank = float(torch.tensor(mean_ranks, dtype=torch.float).mean())
        hits_at_k = int(torch.tensor(hits_at_k).sum()) / len(hits_at_k)

        return mrr, mean_rank, hits_at_k

Thank you for reading!

[rusty1s]

Definitely. Are you interested in contributing this?

[songsong0425]

Sure! However, I'm concerned about making errors during contributions since I'm not good at git-related work. If it is okay, please announce to me how to contribute to it. Thank you for your suggestion! :)

[rusty1s]

I am happy to guide you. Here is how you can do it:

• Fork PyG to your account by clicking the fork button.
• Clone your personal PyG repository via git clone https://github.com/songsong0425/pytorch_geometric
• Checkout a new branch: git checkout -b kge_metric
• Make the changes
• Add your changes to Git: git add . -A
• Commit them: git commit -m "update"
• Push them: git push
• Afterwards, go to PyG GitHub and click "Submit PR"

[songsong0425]

Thank you for your kind explanation!
Please check the request that didn't automatically process by the error..

[rusty1s]

#8298

[sjnov11]

Hi, When can i use this functionallity?

[akihironitta]

@sjnov11 Yes, but you need to install it from the main branch or nightly version: https://github.com/pyg-team/pytorch_geometric#nightly-and-master Your PR #8298 will be included in the next feature release 2.5.0 :)