Repository: pykeen/pykeen
Issue: How do I avoid running out of GPU memory when training on a big dataset?
Root cause: Documentation Deficiency

### Describe the bug

I am trying to use Pykeen pipeline to embed DBpedia (28 million triples). Unfortunately, I encounter insufficient gpu memory error. I tried to fix this by adding model_kwargs={'automatic_memory_optimization': True} in the pipeline. However, the error still remains. I am using TransE.

Is there a way I can specify the batch size in the pipeline directly?

### How to reproduce

```python
for model in args.models:
            print()
            print('*'*50)
            print(f'Training with {model}...')
            print('*'*50)
            result = pipeline(
                training=kg+"/train.txt",
                testing=kg+"/test.txt",
                model=model,
                model_kwargs={'automatic_memory_optimization': True},
                loss=args.loss,
                epochs=args.epochs,
                random_seed=args.random_seed,
                dataset_kwargs=args.dataset_kwargs
            )
            storage_path = kg+'/embeddings/'+'/'+model
            result.save_to_directory(storage_path)
```

### Environment

| Key             | Value                      |
|-----------------|----------------------------|
| OS              | posix                      |
| Platform        | Linux                      |
| Release         | 3.10.0-327.36.1.el7.x86_64 |
| Time            | Thu Jan 20 13:44:45 2022   |
| Python          | 3.9.7                      |
| PyKEEN          | 1.6.0                      |
| PyKEEN Hash     | UNHASHED                   |
| PyKEEN Branch   |                            |
| PyTorch         | 1.9.1                      |
| CUDA Available? | true                       |
| CUDA Version    | 10.2                       |
| cuDNN Version   | 7605                       |  

### Additional information

_No response_

Hi @Jean-KOUAGOU ,

you can set the batch size via the `training_kwargs` parameter (cf. https://pykeen.readthedocs.io/en/stable/api/pykeen.pipeline.pipeline.html#pykeen.pipeline.pipeline)

```python
result = pipeline(
    training_kwargs=dict(batch_size=32),
    ...,
)
```

It may still happen that you encounter a memory error: the TransE model needs to store `(num_entities + num_relations) * embedding_dim` parameters. You can try to reduce the `embedding_dim`, which by default is set to 50, via the `model_kwargs` parameter.

```python
result = pipeline(
    model_kwargs=dict(
        embedding_dim=16,
        ...,
    ),
    ...,
)
```

Thank you very much. I will try your suggestion

I have reduced the batch size to 16 and the embedding dimension to 5, and it still does not run. I encounter the same memory issue

@Jean-KOUAGOU 

what is

* the number of entities,
* the number of relations, and
* your GPU memory size?

There are 6,380,141 entities and 552 relations

GPU memory: vmem usage 87805080k exceeded limit 75497472k on host gpu004

> There are 6, 380, 141 entities and 552 relations

With `embedding_dim=5` this would lead to `4 * embedding_dim * (num_entities + num_relations) = 127,613,860` (the factor of 4 for using the default data type `float32`), i.e., approx. `127MByte` of required memory to store the embeddings, and a negligible amount for the batch of scores. So with your GPU with ~75GiB you should not have any problems :thinking: 

Are you sure you are the only one using the GPU and you do not have another job running on the same GPU?

I am not sure if someone is using the same gpu or not.

Please take a look whether someone else is using the GPU at the same time / whether there is sufficient memory available. You may find one of the following links useful:
* https://pytorch.org/docs/stable/notes/cuda.html#memory-management
* https://pypi.org/project/gpustat/
* https://unix.stackexchange.com/questions/38560/gpu-usage-monitoring-cuda