Repository: pykeen/pykeen
Issue: How do I get the embedding matrix for all entities and relations?
GitHub labels: question

I'm interested in getting the resulting embeddings of entities and relations.
How to get this, please? 
I couldn't find a method which could gives this. 
It seems to me that the `PipelineResult `may have a method for this purpose?

Hi  @xin-zhao0,

`pipeline_result` contains the trained model which contains the embeddings:

```python3
model = pipeline_result.model
entity_embeddings = model.entity_embeddings.weight.data
relation_embeddings = model.relation_embeddings.weight.data
```
Each row in `entity_embeddings/relation_embeddings` contains the embedding of a specific entity/relation. For instance, `entity_embeddings[0]` returns the embedding of the entity with id 0.

To get the ids of the entites/relations, you can use the internal mappings:

```python3
entity_to_id = model.triples_factory.entity_to_id
relation_to_id = model.triples_factory.relation_to_id
```

Putting everything together you can get the entity embeddings for the entities 'brazil' and 'china' contained in the  Nations dataset as follows:

```python3
import torch
import numpy as np

model = pipeline_result.model
entity_embeddings = model.entity_embeddings.weight.data
relation_embeddings  = model.relation_embeddings.weight.data
entity_to_id = model.triples_factory.entity_to_id

entities = ['brazil', 'china']
entity_ids = torch.tensor([entity_to_id[entity] for entity in entities], dtype=torch.long)
embedded_entities = entity_embeddings[entity_ids]
```





@mali-git before we close this, let's add some documentation in the "getting started" tutorial

Sounds good!

Thank you. for models like rescal and some others, there is a relation-specific matrix M. Is that returned in the results as part of the model?

> Thank you. for models like rescal and some others, there is a relation-specific matrix M. Is that returned in the results as part of the model?

The relations matrices are also saved in `model.relation_embeddings.weight.data`, but they need to be reshaped. For RESCAL, you can get the relation embedding dimension as follows (please consider that some models have a dedicated `relation_dim` which needs to be used when reshaping the relation-embeddings) :

```python3
import torch
import numpy as np

# Get the relation-matrix of the relation with the ID 0
relation_embeddings = rescal.relation_embeddings.weight.data
relation = relation_embeddings[0].view(rescal.embedding_dim, rescal.embedding_dim)

# Get the embeddings for the relations 'officialvisits' and 'reldiplomacy' contaiend in the Nations dataset
relation_to_id = model.triples_factory.relation_to_id
relations = ['officialvisits', 'reldiplomacy']
relation_ids = torch.tensor([relation_to_id[relation] for relation in relations], dtype=torch.long)
relation_embeddings = relation_embeddings[relation_ids].view(-1, rescal.embedding_dim, rescal.embedding_dim)
```

@mali-git to access the weight data, I have to call 

```python
from pykeen.pipeline import pipeline
result = pipeline(
    model='TransE',
    dataset='umls',
)
model = result.model
entity_embeddings = model.entity_embeddings._embeddings.weight.data
relation_embeddings = model.relation_embeddings._embeddings.weight.data
```

Has this been updated recently or did I overlook something? :smile:



> @mali-git to access the weight data, I have to call
> 
> ```python
> from pykeen.pipeline import pipeline
> result = pipeline(
>     model='TransE',
>     dataset='umls',
> )
> model = result.model
> entity_embeddings = model.entity_embeddings._embeddings.weight.data
> relation_embeddings = model.relation_embeddings._embeddings.weight.data
> ```
> 
> Has this been updated recently or did I overlook something? smile

Hi, this way should still be possible.

For a more future-proof way you can use

```python
entity_embeddings = model.entity_embeddings()
```
or more explicitly
```python
entity_embeddings = model.entity_embeddings(indices=None)
```

which does not depend on the underlying Embedding implementation details, but uses the interface [`RepresentationModule`](https://github.com/pykeen/pykeen/blob/68af86ea387cd3282a7561884aee96d97e63dbda/src/pykeen/nn/emb.py#L34).

Also, we are extending our framework to have a unified way to store more than one representation for entities/relations, cf. here:
https://github.com/pykeen/pykeen/blob/68af86ea387cd3282a7561884aee96d97e63dbda/src/pykeen/models/nbase.py#L310

Briefly, the change is to always store a list. In that case, for TransE, you would need to use `model.entity_representations[0]()`. We'll update the docs once we actually did migrate the models.

@shimst3r @xin-zhao0 I've started a PR where we translated some of this discussion into the documentation. Thanks for keeping us honest ;)

New docs at https://pykeen.readthedocs.io/en/latest/tutorial/first_steps.html#using-learned-embeddings

Thanks @cthoyt, very much appreciated. :+1: