Repository: pykeen/pykeen
Issue: How to select which edges to train in heterogeneous graph
Root cause: Documentation Deficiency

### What is your question

Hi, in a knowledge graph many times there are multiple relations, but predicting only some specific type (e.g. gene-disease) is desirable. I have a graph with over 1MM edges, but I am interested in only 10k of them.

I couldn't figure out how add to the pipeline my whole dataset, with all edges, but only have the model (RGCN) train a specific edge type. Any ideas would be appreciated. Thanks!

Felipe

### Environment

nan

### Issue Template Checks

- [X] This is not a bug report (use a different issue template if it is)
- [X] This is not a feature request (use a different issue template if it is)
- [X] I've read the text explaining why including environment information is important and understand if I omit this information that my issue will be dismissed

Hi @fmellomascarenhas, I'd suggest carefully reading through the documentation and other similar issues in this repo, such as:

- https://pykeen.readthedocs.io/en/stable/tutorial/making_predictions.html
- https://github.com/pykeen/pykeen/issues/863
- https://github.com/pykeen/pykeen/issues/337
- https://github.com/pykeen/pykeen/issues/84

Are you looking into this in a capacity as an industrial research scientist?

Hi @cthoyt , I am sorry, I should have expressed myself better. I understand how to evaluate the model on a specific edge set. I believe that my questions is related to **training**.

I want to minimize the loss of the 10k DRUG-DISEASE edge pairs in the data, and not all 1 million edge pairs that exist in the whole dataset.

My question does not make sense for an algorithm like TransE, because If I want to optimize only DRUG-DISEASE, I can provide that alone. However, in RGCN, I need to have my whole dataset, because the neighborhood information is important. So is there a way to provide the whole dataset, but only optimize the loss for a specific set of relations?

Please, let me know if that makes sense, and thanks for taking the time to answer. Answering your question, yes, I work for a company as a data scientist. Thanks! :)



Hi @fmellomascarenhas, thanks for making it a bit more clear. It's for sure an interesting question that has a potential answer, but right now our team isn't yet sure about the best way to respond to industrial requests for support outside of the context of paid consulting.

In the mean time, you can sit tight and we might be able to focus on this if there's wide community interest, you could take a look at the training loop itself and get a feel for how the training data loaders, the loss, and the loop itself interact by checking out the following code, or you could share your email so we could get in touch privately about alternative ways to go forward.

https://github.com/pykeen/pykeen/blob/e49a3551f4a54ea246a5233ce6d0aa0291f1c313/src/pykeen/training/training_loop.py#L603-L704

Don't have an MWP at hand, but conceptually it should be easy:
1. Initializing the GNN encoder with the full graph from the whole set of training edges
2. Limiting the training set (sent to the training loop over which the loss will be computed) to only those desired 10k, can be done in 1-2 lines with edge masking by relation type

This way, the forward pass will consider all edges in a graph, but the optimization will happen only for the 10k edges

Thanks @migalkin and @cthoyt. I will investigate it and let you know how it went. Thank you very much!