Repository: pykeen/pykeen
Issue: Out of memory after some epochs, for some models
Root cause: Resource Mishandling

### Describe the bug

Hi. I am training different PyKEEN models (TransE, DistMult,..) using them directly (without the pipeline). There is a strange behaviour. For example, when I use TransE, everything goes fine. When I change to DistMult, the memory increases every epoch and after some epochs I receive a OOM error. 

### How to reproduce

The code is the following:

```python
from pykeen.models import TransE, DistMult
from pykeen.triples import TriplesFactory
import pandas as pd
import torch as th
import torch.optim as optim
import torch.nn as nn
from tqdm import tqdm, trange

class FastTensorDataLoader:
    """
    A DataLoader-like object for a set of tensors that can be much faster than
    TensorDataset + DataLoader because dataloader grabs individual indices of
    the dataset and calls cat (slow).
    Source: https://discuss.pytorch.org/t/dataloader-much-slower-than-manual-batching/27014/6
    """

    def __init__(self, *tensors, batch_size=32, shuffle=False):
        """
        Initialize a FastTensorDataLoader.
        :param *tensors: tensors to store. All tensors must have the same size at dimension 0.
        :param batch_size: batch size to load. Defaults to 32.
        :type batch_size: int, optional
        :param shuffle: if True, shuffle the data *in-place* whenever an
            iterator is created out of this object. Defaults to False.
        :type shuffle: bool, optional
        """

        # Type checking
        if not all(isinstance(t, th.Tensor) for t in tensors):
            raise TypeError("All non-optional parameters must be Tensors")

        if not isinstance(batch_size, int):
            raise TypeError("Optional parameter batch_size must be of type int")

        if not isinstance(shuffle, bool):
            raise TypeError("Optional parameter shuffle must be of type bool")

        assert all(t.shape[0] == tensors[0].shape[0] for t in tensors)
        self.tensors = tensors

        self.dataset_len = self.tensors[0].shape[0]
        self.batch_size = batch_size
        self.shuffle = shuffle

        # Calculate # batches
        n_batches, remainder = divmod(self.dataset_len, self.batch_size)
        if remainder > 0:
            n_batches += 1
        self.n_batches = n_batches

    def __iter__(self):
        if self.shuffle:
            r = th.randperm(self.dataset_len)
            self.tensors = [t[r] for t in self.tensors]
        self.i = 0
        return self

    def __next__(self):
        if self.i >= self.dataset_len:
            raise StopIteration
        batch = tuple(t[self.i:self.i + self.batch_size] for t in self.tensors)
        self.i += self.batch_size
        return batch

    def __len__(self):
        return self.n_batches

    
graph_file = 'path/to/some/tsv/file'
graph_df = pd.read_csv(graph_file, sep='\t', header=None)
graph_df.columns = ['head', 'relation', 'tail']
entities = graph_df['head'].append(graph_df['tail']).unique()
relations = graph_df['relation'].unique()
entities.sort()
relations.sort()
entity_to_id = {entity: i for i, entity in enumerate(entities)}
relation_to_id = {relation: i for i, relation in enumerate(relations)}

heads = [entity_to_id[head] for head in graph_df['head']]
tails = [entity_to_id[tail] for tail in graph_df['tail']]
relations = [relation_to_id[relation] for relation in graph_df['relation']]

triples = list(zip(heads, relations, tails))
triples = th.tensor(triples, dtype=th.long)
assert triples.shape[1] == 3, f'Expected triples of shape (n,3), got {triples.shape}'
triples_factory = TriplesFactory(triples, entity_to_id, relation_to_id)

## HERE I CHANGE THE MODEL
model = "distmult"
if model == "transe":
    model = TransE(triples_factory=triples_factory, embedding_dim=1024)
elif model == "distmult":
    model = DistMult(triples_factory=triples_factory, embedding_dim=1024)
    
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion_bpr = nn.LogSigmoid()

device = "cuda"
model = model.to(device)

train_dataloader = FastTensorDataLoader(triples[:,0], triples[:,1], triples[:,2], batch_size=128, shuffle=True)

best_loss = float("inf")
best_mr = float("inf")

num_negs = 2
for epoch in trange(100, desc=f"Training..."):
    model.train()

    graph_loss = 0
    for head, rel, tail in tqdm(train_dataloader, desc="Processing batches"):
        head = head.to(device)
        rel = rel.to(device)
        tail = tail.to(device)
                
        pos_logits = model.forward(head, rel, tail, mode=None)

        neg_logits = 0
        for i in range(num_negs):
            neg_tail = th.randint(0, 100, (len(head),), device=device)
            neg_logits += model.forward(head, rel, neg_tail, mode=None)
                                                            
        neg_logits /= (num_negs)

        batch_loss = -criterion_bpr(pos_logits - neg_logits - 0.01).mean()
                
        optimizer.zero_grad()
        batch_loss.backward()
        optimizer.step()
                
        graph_loss += batch_loss.item()
        graph_loss /= len(train_dataloader)

```                                                                                                                                                                                                                         

The error I get with DistMult is:

```python-traceback
Traceback (most recent call last):
  File "example.py", line 129, in <module>
    optimizer.step()
  File "/home/zhapacfp/miniconda3/envs/semrec/lib/python3.8/site-packages/torch/optim/optimizer.py", line 280, in wrapper
    out = func(*args, **kwargs)
  File "/home/zhapacfp/miniconda3/envs/semrec/lib/python3.8/site-packages/torch/optim/optimizer.py", line 33, in _use_grad
    ret = func(self, *args, **kwargs)
  File "/home/zhapacfp/miniconda3/envs/semrec/lib/python3.8/site-packages/torch/optim/adam.py", line 141, in step
    adam(
  File "/home/zhapacfp/miniconda3/envs/semrec/lib/python3.8/site-packages/torch/optim/adam.py", line 281, in adam
    func(params,
  File "/home/zhapacfp/miniconda3/envs/semrec/lib/python3.8/site-packages/torch/optim/adam.py", line 507, in _multi_tensor_adam
    denom = torch._foreach_add(exp_avg_sq_sqrt, eps)
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 38.00 MiB (GPU 0; 11.90 GiB total capacity; 11.68 GiB already allocated; 36.12 MiB free; 11.71 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

```


### Environment

| Key             | Value                    |
|-----------------|--------------------------|
| OS              | posix                    |
| Platform        | Linux                    |
| Release         | 5.4.0-152-generic        |
| Time            | Wed Jul 19 02:14:35 2023 |
| Python          | 3.8.0                    |
| PyKEEN          | 1.10.1                   |
| PyKEEN Hash     | UNHASHED                 |
| PyKEEN Branch   |                          |
| PyTorch         | 2.0.1+cu117              |
| CUDA Available? | true                     |
| CUDA Version    | 11.7                     |
| cuDNN Version   | 8500                     |

### Additional information

_No response_

### Issue Template Checks

- [X] This is not a feature request (use a different issue template if it is)
- [X] This is not a question (use the [discussions](https://github.com/pykeen/pykeen/discussions/new) forum instead)
- [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

+1. I have the same issue when I used the forward() function and loss() function seperately in the trainer. 

Hi @ferzcam (and @renzhonglu11 ),

this likely comes `DistMult` using a regularizer by default, cf.
https://github.com/pykeen/pykeen/blob/d1222b7c18d494290d285525b76c3f22e30db467/src/pykeen/models/unimodal/distmult.py#L85-L86 that needs to accumulate the regularization term across micro-batches. PyKEEN's pipelines call `collect_regularization_term` which not only collects the regularization terms across different places, but also releases their references. When this is not called, the term accumulates indefinetely, and holds references to buffers -> over time you will run out of memory.

To fix it, either

- configure the model without regularizer, e.g., for `DistMult` pass `regularizer=None`, or
- if you want to make use of the regularizer, make sure to call `collect_regularization_term` on the models and include this into your loss terms

Thanks a lot. But I am still quite not sure if it is really the reason. I took a look at Pykeen's source code. It seems Pykeen calculates the loss just with a function in the trainer. (https://github.com/pykeen/pykeen/blob/d1222b7c18d494290d285525b76c3f22e30db467/src/pykeen/training/training_loop.py#L643)
What I did is that I first call forward() of a model(like Distmult) to calculate the prediction value and then call loss to calculate the loss value. (similar to what  @ferzcam did) Then the memory will increase per epoch during training despite using GPU or CPU.
However, when I put forward and loss calculation together in one function just like Pykeen, and call this function. The memory does not increase anymore. So confused about it. 🧐
***
In my KGE Model
```python
class PykeenKGE:
    def training_step(self, batch):
        x_batch, y_batch = batch
        yhat_batch = self.forward(x_batch)
        loss_batch = self.loss(yhat_batch, y_batch)
        return loss_batch + self.model.collect_regularization_term()
```
In Trainer: 
```python
 batch_loss = self.model.training_step(batch)
```


@mberr your solution works. Thanks a lot!😁

You can see with

```python
from pykeen.datasets import get_dataset
from pykeen.models import DistMult

dataset = get_dataset(dataset="nations")
model = DistMult(triples_factory=dataset.training)
print(model)
```

how the resulting structure looks like:
```
DistMult(
  (loss): MarginRankingLoss(
    (margin_activation): ReLU()
  )
  (interaction): DistMultInteraction()
  (entity_representations): ModuleList(
    (0): Embedding(
      (_embeddings): Embedding(14, 50)
    )
  )
  (relation_representations): ModuleList(
    (0): Embedding(
      (regularizer): LpRegularizer()
      (_embeddings): Embedding(55, 50)
    )
  )
  (weight_regularizers): ModuleList()
)
```

Notice how the `relation_representations` as a `regularizer`. compare this to
```python
from pykeen.datasets import get_dataset
from pykeen.models import DistMult

dataset = get_dataset(dataset="nations")
model = DistMult(triples_factory=dataset.training, regularizer=None)
print(model)

```
resulting in
```
DistMult(
  ...
  (relation_representations): ModuleList(
    (0): Embedding(
      (_embeddings): Embedding(55, 50)
    )
  )
  ...
)
```

The regularization term of the relation embedding is updated here
https://github.com/pykeen/pykeen/blob/d1222b7c18d494290d285525b76c3f22e30db467/src/pykeen/nn/representation.py#L189-L191
i.e., in the `Embedding`'s forward call.

Nice. Now it makes more sense.  And I also found out where Pykeen calls the `collect_regularization_term()`. Thanks for your explanation.

Hi @mberr. Thanks for the explanation. I was able to make it work now!