Repository: pykeen/pykeen
Issue: Load pre-trained embeddings
Root cause: Documentation Deficiency

Hi,
How can I initialize my entities with pre-trained embeddings obtained from an external task? Just for entities embeddings, but not for relation embeddings.
Thanks,

Hi @alipsgh ,

the easiest way would be to use torch's `load_state_dict` (since every pykeen Model is also a `torch.nn.Module`). For this, you need to get the name of the entity embedding weight, and prepare a dictionary which maps this name to the pretrained embedding.

Here is an example:
```python
import torch
from pykeen.models import TransE
from pykeen.datasets import get_dataset

dataset = get_dataset(dataset="nations")
model = TransE(
    triples_factory=dataset.training,
    embedding_dim=128,  # has to match dimension of pre-trained embedding
)
print(model.state_dict())   # use this to find the name of the weight

pretrained_embedding_tensor = torch.rand(dataset.num_entities, 128)  # replace this by torch.load(pretrained_embeddings_path) or similar

model.load_state_dict(state_dict={
    "entity_embeddings._embeddings.weight": pretrained_embedding_tensor,
}, strict=False)
```

You can find more general information about `torch.load_state_dict` [here](https://pytorch.org/tutorials/beginner/saving_loading_models.html#warmstarting-model-using-parameters-from-a-different-model)

EDIT: here is another example for how to use the `entity_initializer` keyword to achieve the same for the high-level `pipeline` function. This method may only work for a subset of models and/or embeddings. For models with more than one weight per entity (e.g. `TransH`), you may need to use the first method.

```python
import torch
from pykeen.pipeline import pipeline


num_entities = 14
pretrained_embedding_tensor = torch.rand(num_entities, 128)


def initialize_from_pretrained(x: torch.FloatTensor) -> torch.FloatTensor:
    return pretrained_embedding_tensor


result = pipeline(
    dataset="nations",
    model="transe",
    model_kwargs=dict(
        embedding_dim=pretrained_embedding_tensor.shape[-1],
        entity_initializer=initialize_from_pretrained,
    ),
)
```

Update: I'll add a small utility method to simplify the second case in #638 

See also https://github.com/pykeen/pykeen/issues/356

Hi @alipsgh, we just provided a technical solution for this and a short tutorial in https://github.com/pykeen/pykeen/pull/638

Thank you, @mberr and @cthoyt  🙂 

I tried your suggestions for my pre-trained embeddings, unfortunately, none of them worked.

## RE: First suggestion

I think the embeddings are reset by the `reset_parameters` function somewhere in the pipeline. To avoid any reset, I had to change the `reset_parameters` and `post_parameter_update` functions by adding the `return` statement before other statements in `emb.py` -  I however still doubt that they are not changed somewhere else.

```python
    def reset_parameters(self) -> None:  # noqa: D102
        return
        # initialize weights in-place
        self._embeddings.weight.data = self.initializer(
            self._embeddings.weight.data.view(self.num_embeddings, *self.shape),
        ).view(self.num_embeddings, self.embedding_dim)
        
    def post_parameter_update(self):  # noqa: D102
        return
        # apply constraints in-place
        if self.constrainer is not None:
            self._embeddings.weight.data = self.constrainer(self._embeddings.weight.data)
```

## RE: Second Suggestion

Regarding the second suggestion, which I found more elegant, it seems that the `create_init_from_pretrained` function doesn't work properly - particularly when it is set within a pipeline. Here's what I did:

1. I set `epochs` to `0`
2. After the run, I checked the model's **entity weights**, and they are different from the pre-trained ones. (There was no training since I set the num. of epochs to `0`.)
- For the resolved first suggestion, the embeddings didn't change in the pipeline when `epochs=0`.

Hi @mberr 

1️⃣ 
I debugged the code and changed this part of the code in the `training_loop.py` file:
```python
    # when called by batch_size_search(), the parameter update should not be applied.
    if not only_size_probing:
        callback.pre_step()

        # update parameters according to optimizer
        self.optimizer.step()

        # After changing applying the gradients to the embeddings, the model is notified that the forward
        # constraints are no longer applied
        self.model.post_parameter_update()
```
I had to move `self.model.post_parameter_update()` into the `if` statement to fix the issue. It was outside the if statement.

2️⃣ 
❓ I have another question! How can I use your solution within an HPO pipeline if I want to pass pre-trained embeddings with different sizes as a dict to `entity_initializer`?

@alipsgh thank you for looking into it! Would you mind to open a PR with the changes you suggest?

Regarding your second point:
Do I understand you correctly, that you have a set of different pre-trained embeddings, which have different embedding dimensions?



Hi @mberr , thanks for your response. @alipsgh and I work on the same project
Yes, we have different pre-trained embeddings of different sizes, let's say 32, 64, 128, 196, and 256, and we want HPO to suggest the best one when a model is initialized with them. Alongside searching for the best pre-trained embeddings we should be able to consider other ranges for other arguments in `model_kwargs_ranges` too, such as `entity_constrainer=dict(type='categorical', choices=[CHOICE1, CHOICE2, etc.])`) 

Max (@mberr) I'll open a PR when we run multiple experiments to ensure it is a safe change.
As to your question, Yes! Just like what Ahmad mentioned above.

Hi @mberr and @cthoyt 

I noticed that there's been a migration from the [create_init_from_pretrained](https://github.com/pykeen/pykeen/pull/638#issue-1053387266) function to the [PretrainedInitializer](https://github.com/pykeen/pykeen/blob/e401e96c4400ccdc2ff7d4f7c379f16829f6fe82/src/pykeen/nn/init.py#L128) class.
I checked the class for a few runs and it worked without any issue.

I am going to close this issue.

Thanks 💯