Repository: pykeen/pykeen
Issue: Pykeen in OGBL-Biokg: #Node and out of RAM
Root cause: Resource Mishandling

### Describe the bug

Hi, I am currently running for OGBL-Biokg using pykeen, and the total number of node from the ogbl-biokg should be around 90000 according to the original paper as well as ogb. However, in pykeen documentation there seems to be only around 45000 nodes. Is there any clarification on why is it the case?

In addition, I am trying to run RGCN in OGBL-Biokg on pykeen, but I constantly having out-of-RAM error (it requires more than 80GB to load dataset). I am using 80GB ram with 40GB A100. Is there any fix for it?



### How to reproduce

```python
import pykeen
from pykeen.datasets import OGBBioKG
from pykeen.pipeline import pipeline

dataset = OGBBioKG()
model = "RGCN"
metrics = ["hits@1","hits@3","hits@10", "mean_rank", "mean_reciprocal_rank"]

# Define the training configuration
training_kwargs = dict(
    num_epochs=10,  # Number of training epochs
    batch_size=128,  # Batch size for training
)
result = pipeline(
    dataset=dataset,
    model=model,
    evaluation_relation_whitelist=None,  # Evaluate on all relations
    training_kwargs=training_kwargs,
    model_kwargs=dict(embedding_dim=100),  # Set the embedding dimension
    loss="BCEWithLogitsLoss",  # Use binary cross-entropy loss with logits
    loss_kwargs=dict(),  # No additional loss parameters
    regularizer=None,  # No regularizer
    optimizer="Adam",  # Use the Adam optimizer
    optimizer_kwargs=dict(lr=0.00001),  # Set the learning rate
    evaluator="RankBasedEvaluator",  # Use the rank-based evaluator
    evaluator_kwargs=dict(filtered=True),  # Use filtered evaluation
    evaluation_kwargs=dict(batch_size=256),  # Batch size for evaluation
)
```

### Environment

```
| Key             | Value                    |
|-----------------|--------------------------|
| OS              | posix                    |
| Platform        | Linux                    |
| Release         | 5.10.147+                |
| Time            | Thu Mar 30 16:39:44 2023 |
| Python          | 3.9.16                   |
| PyKEEN          | 1.10.1                   |
| PyKEEN Hash     | UNHASHED                 |
| PyKEEN Branch   |                          |
| PyTorch         | 1.13.1+cu116             |
| CUDA Available? | true                     |
| CUDA Version    | 11.6                     |
| cuDNN Version   | 8302                     |
```

### 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

> Hi, I am currently running for OGBL-Biokg using pykeen, and the total number of node from the ogbl-biokg should be around 90000 according to the original paper as well as ogb. However, in pykeen documentation there seems to be only around 45000 nodes. Is there any clarification on why is it the case?

It looks like the number of nodes is wrongly annotated here
https://github.com/pykeen/pykeen/blob/b3ef9d19883967b41f7fdc2b540972ebcb35f273/src/pykeen/datasets/ogb.py#L98-L104

which then gets rendered for the documentation. Maybe this has been wrongly transferred from the paper, since the number `45085` is actually mentioned in the paper - as the number of entities of type protein functions alone:

> It contains 5 types of entities: diseases (10,687 nodes), proteins (17,499), drugs (10,533 nodes), side effects (9,969 nodes), and protein functions (45,085 nodes) [[Source]](https://arxiv.org/pdf/2005.00687v7.pdf)

> In addition, I am trying to run RGCN in OGBL-Biokg on pykeen, but I constantly having out-of-RAM error (it requires more than 80GB to load dataset). I am using 80GB ram with 40GB A100. Is there any fix for it?

Could you provide a traceback / concrete error messages?

> > Hi, I am currently running for OGBL-Biokg using pykeen, and the total number of node from the ogbl-biokg should be around 90000 according to the original paper as well as ogb. However, in pykeen documentation there seems to be only around 45000 nodes. Is there any clarification on why is it the case?
> 
> It looks like the number of nodes is wrongly annotated here
> 
> https://github.com/pykeen/pykeen/blob/b3ef9d19883967b41f7fdc2b540972ebcb35f273/src/pykeen/datasets/ogb.py#L98-L104
> 
> which then gets rendered for the documentation. Maybe this has been wrongly transferred from the paper, since the number `45085` is actually mentioned in the paper - as the number of entities of type protein functions alone:
> 
> > It contains 5 types of entities: diseases (10,687 nodes), proteins (17,499), drugs (10,533 nodes), side effects (9,969 nodes), and protein functions (45,085 nodes) [[Source]](https://arxiv.org/pdf/2005.00687v7.pdf)

yes, and I believe this has affected a lot of recent papers that tested OGBBiokg based on pykeen, as from what I can search, a lot of paper that use pykeen for their experiments on OGBBioKG mentioned 45,085 in their description of datasets. 

> > In addition, I am trying to run RGCN in OGBL-Biokg on pykeen, but I constantly having out-of-RAM error (it requires more than 80GB to load dataset). I am using 80GB ram with 40GB A100. Is there any fix for it?
> 
> Could you provide a traceback / concrete error messages?

Yes. In the colab, I can see my usage of RAM constantly goes up and then when it hits the limit, the session crashes.
By the time system crash, we have the following output from console:
```
WARNING:pykeen.pipeline.api:No random seed is specified. Setting to 828588355.
Downloading http://snap.stanford.edu/ogb/data/linkproppred/biokg.zip
Downloaded 0.90 GB: 100%|██████████| 920/920 [00:18<00:00, 48.65it/s]
Extracting /root/.data/pykeen/datasets/ogbbiokg/biokg.zip
Loading necessary files...
This might take a while.
Processing graphs...
100%|██████████| 1/1 [00:00<00:00, 4112.06it/s]Saving...
```

On the slurm cluster server, the same happened with message
```
slurmstepd: error: Detected 1 oom-kill event(s) in StepId=2597285.batch cgroup. Some of your processes may have been killed by the cgroup out-of-memory handler.
```
 

> Yes. In the colab, I can see my usage of RAM constantly goes up and then when it hits the limit, the session crashes.

I took a look at the code and might have found the reason. Could you try whether the branch `improve-ogb-loading` fixes your problem? The following should work

```console
pip install https://github.com/pykeen/pykeen/archive/improve-ogb-loading.zip
```

You may need to additionally provide `--force-reinstall`

> > Yes. In the colab, I can see my usage of RAM constantly goes up and then when it hits the limit, the session crashes.
> 
> I took a look at the code and might have found the reason. Could you try whether the branch `improve-ogb-loading` fixes your problem? The following should work
> 
> ```
> pip install https://github.com/pykeen/pykeen/archive/improve-ogb-loading.zip
> ```
> 
> You may need to additionally provide `--force-reinstall`

Thank you for your prompt fix! Unfortunately, another error occur when I run this:
```
WARNING:pykeen.pipeline.api:No random seed is specified. Setting to 2708553551.
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
[<ipython-input-10-dcb5236d43bd>](https://localhost:8080/#) in <cell line: 13>()
     11 # RGCN internally creates inverse triples. It thus expects a triples factory without them.
     12 # Run the pipeline
---> 13 result = pipeline(
     14     dataset=dataset,
     15     model=model,

8 frames
[/usr/local/lib/python3.9/dist-packages/pykeen/pipeline/api.py](https://localhost:8080/#) in pipeline(dataset, dataset_kwargs, training, testing, validation, evaluation_entity_whitelist, evaluation_relation_whitelist, model, model_kwargs, interaction, interaction_kwargs, dimensions, loss, loss_kwargs, regularizer, regularizer_kwargs, optimizer, optimizer_kwargs, clear_optimizer, lr_scheduler, lr_scheduler_kwargs, training_loop, training_loop_kwargs, negative_sampler, negative_sampler_kwargs, epochs, training_kwargs, stopper, stopper_kwargs, evaluator, evaluator_kwargs, evaluation_kwargs, result_tracker, result_tracker_kwargs, metadata, device, random_seed, use_testing_data, evaluation_fallback, filter_validation_when_testing, use_tqdm)
   1496     _result_tracker.start_run(run_name=title)
   1497 
-> 1498     training, testing, validation = _handle_dataset(
   1499         _result_tracker=_result_tracker,
   1500         dataset=dataset,

[/usr/local/lib/python3.9/dist-packages/pykeen/pipeline/api.py](https://localhost:8080/#) in _handle_dataset(_result_tracker, dataset, dataset_kwargs, training, testing, validation, evaluation_entity_whitelist, evaluation_relation_whitelist)
    921         )
    922 
--> 923     training, testing, validation = dataset_instance.training, dataset_instance.testing, dataset_instance.validation
    924     # evaluation restriction to a subset of entities/relations
    925     if any(f is not None for f in (evaluation_entity_whitelist, evaluation_relation_whitelist)):

[/usr/local/lib/python3.9/dist-packages/pykeen/datasets/base.py](https://localhost:8080/#) in training(self)
    332         """The training triples factory."""
    333         if not self._loaded:
--> 334             self._load()
    335         assert self._training is not None
    336         return self._training

[/usr/local/lib/python3.9/dist-packages/pykeen/datasets/ogb.py](https://localhost:8080/#) in _load(self)
     51         dataset = LinkPropPredDataset(name=self.name, root=self.cache_root)
     52         edge_split = dataset.get_edge_split()
---> 53         self._training = self._make_tf(edge_split["train"])
     54         assert self._training is not None  # makes mypy hapy
     55         self._testing = self._make_tf(

[/usr/local/lib/python3.9/dist-packages/pykeen/datasets/ogb.py](https://localhost:8080/#) in _make_tf(self, x, entity_to_id, relation_to_id)
    121         # note: to be more memory-efficient we do not convert to numpy array of strings
    122         # heads and tails need to be prefixed by the type to avoid name collision
--> 123         return super()._make_tf(
    124             x=dict(
    125                 relation=x["relation"],

[/usr/local/lib/python3.9/dist-packages/pykeen/datasets/ogb.py](https://localhost:8080/#) in _make_tf(self, x, entity_to_id, relation_to_id)
     75         # note: to be more memory-efficient we do not convert to numpy array of strings
     76         entity_to_id = entity_to_id or create_entity_mapping(heads=x["head"], tails=x["tail"])
---> 77         relation_to_id = relation_to_id or create_relation_mapping(relations=x["relation"])
     78         # convert to integers
     79         mapped_triples = cast(

[/usr/local/lib/python3.9/dist-packages/pykeen/triples/triples_factory.py](https://localhost:8080/#) in create_relation_mapping(relations)
    107     """
    108     # Sorting ensures consistent results when the triples are permuted
--> 109     relation_labels = sorted(
    110         set(relations),
    111         key=lambda x: (re.sub(f"{INVERSE_SUFFIX}$", "", x), x.endswith(f"{INVERSE_SUFFIX}")),

[/usr/local/lib/python3.9/dist-packages/pykeen/triples/triples_factory.py](https://localhost:8080/#) in <lambda>(x)
    109     relation_labels = sorted(
    110         set(relations),
--> 111         key=lambda x: (re.sub(f"{INVERSE_SUFFIX}$", "", x), x.endswith(f"{INVERSE_SUFFIX}")),
    112     )
    113     # Create mapping

[/usr/lib/python3.9/re.py](https://localhost:8080/#) in sub(pattern, repl, string, count, flags)
    208     a callable, it's passed the Match object and must return
    209     a replacement string to be used."""
--> 210     return _compile(pattern, flags).sub(repl, string, count)
    211 
    212 def subn(pattern, repl, string, count=0, flags=0):

TypeError: cannot use a string pattern on a bytes-like object
```

I have re-tested it today, and the aforementioned error message disappear but replaced with another one. I thought pasting here might be helpful
```
KeyError    Traceback (most recent call last)
[<ipython-input-7-dcb5236d43bd>](https://localhost:8080/#) in <cell line: 13>()
     11 # RGCN internally creates inverse triples. It thus expects a triples factory without them.
     12 # Run the pipeline
---> 13 result = pipeline(
     14     dataset=dataset,
     15     model=model,

6 frames
[/usr/local/lib/python3.9/dist-packages/pykeen/pipeline/api.py](https://localhost:8080/#) in pipeline(dataset, dataset_kwargs, training, testing, validation, evaluation_entity_whitelist, evaluation_relation_whitelist, model, model_kwargs, interaction, interaction_kwargs, dimensions, loss, loss_kwargs, regularizer, regularizer_kwargs, optimizer, optimizer_kwargs, clear_optimizer, lr_scheduler, lr_scheduler_kwargs, training_loop, training_loop_kwargs, negative_sampler, negative_sampler_kwargs, epochs, training_kwargs, stopper, stopper_kwargs, evaluator, evaluator_kwargs, evaluation_kwargs, result_tracker, result_tracker_kwargs, metadata, device, random_seed, use_testing_data, evaluation_fallback, filter_validation_when_testing, use_tqdm)
   1496     _result_tracker.start_run(run_name=title)
   1497 
-> 1498     training, testing, validation = _handle_dataset(
   1499         _result_tracker=_result_tracker,
   1500         dataset=dataset,

[/usr/local/lib/python3.9/dist-packages/pykeen/pipeline/api.py](https://localhost:8080/#) in _handle_dataset(_result_tracker, dataset, dataset_kwargs, training, testing, validation, evaluation_entity_whitelist, evaluation_relation_whitelist)
    921         )
    922 
--> 923     training, testing, validation = dataset_instance.training, dataset_instance.testing, dataset_instance.validation
    924     # evaluation restriction to a subset of entities/relations
    925     if any(f is not None for f in (evaluation_entity_whitelist, evaluation_relation_whitelist)):

[/usr/local/lib/python3.9/dist-packages/pykeen/datasets/base.py](https://localhost:8080/#) in training(self)
    332         """The training triples factory."""
    333         if not self._loaded:
--> 334             self._load()
    335         assert self._training is not None
    336         return self._training

[/usr/local/lib/python3.9/dist-packages/pykeen/datasets/ogb.py](https://localhost:8080/#) in _load(self)
     51         dataset = LinkPropPredDataset(name=self.name, root=self.cache_root)
     52         edge_split = dataset.get_edge_split()
---> 53         self._training = self._make_tf(edge_split["train"])
     54         assert self._training is not None  # makes mypy hapy
     55         self._testing = self._make_tf(

[/usr/local/lib/python3.9/dist-packages/pykeen/datasets/ogb.py](https://localhost:8080/#) in _make_tf(self, x, entity_to_id, relation_to_id)
    121         # note: to be more memory-efficient we do not convert to numpy array of strings
    122         # heads and tails need to be prefixed by the type to avoid name collision
--> 123         return super()._make_tf(
    124             x=dict(
    125                 relation=x["relation"],

[/usr/local/lib/python3.9/dist-packages/pykeen/datasets/ogb.py](https://localhost:8080/#) in _make_tf(self, x, entity_to_id, relation_to_id)
     80             MappedTriples,
     81             torch.as_tensor(
---> 82                 data=[
     83                     [entity_to_id[h], relation_to_id[r], entity_to_id[t]]
     84                     for h, r, t in zip(x["head"], x["relation"], x["tail"])

[/usr/local/lib/python3.9/dist-packages/pykeen/datasets/ogb.py](https://localhost:8080/#) in <listcomp>(.0)
     81             torch.as_tensor(
     82                 data=[
---> 83                     [entity_to_id[h], relation_to_id[r], entity_to_id[t]]
     84                     for h, r, t in zip(x["head"], x["relation"], x["tail"])
     85                 ],

KeyError: 0
```

I can now successfully load the `OGBBioKG` dataset:
```python-console
>>> from pykeen.datasets import OGBBioKG
>>> OGBBioKG().summarize()
OGBBioKG (create_inverse_triples=False)
Name        Entities    Relations      Triples
----------  ----------  -----------  ---------
Training    93773       51             4762678
Testing     93773       51              162870
Validation  93773       51              162886
Total       -           -              5088434
Head              Relation         tail
----------------  ---------------  ------------
disease:C0038586  disease-protein  protein:1653
disease:C0751849  disease-protein  protein:718
disease:C1320474  disease-protein  protein:8622
disease:C0270844  disease-protein  protein:3569
disease:C4279912  disease-protein  protein:8856
```

I still get an OOM error for `WikiKG2`, but his may also be due to the excessive size of this dataset.

Thanks! It is working well in my test for RGCN!    

I still run into another problem in CompGCN on OGBBioKG though:
```
dataset = OGBBioKG()
model = "CompGCN"
metrics = ["hits@1","hits@3","hits@10", "mean_rank", "mean_reciprocal_rank"]

# Define the training configuration
training_kwargs = dict(
    num_epochs=5,  # Number of training epochs
    batch_size=256,  # Batch size for training
)
# Run the pipeline
result = pipeline(
    dataset=dataset,
    model=model,
    evaluation_relation_whitelist=None,  # Evaluate on all relations
    training_kwargs=training_kwargs,
    dataset_kwargs=dict(create_inverse_triples=True), 
    model_kwargs=dict(embedding_dim=100),  # Set the embedding dimension
    loss="BCEWithLogitsLoss",  # Use binary cross-entropy loss with logits
    loss_kwargs=dict(),  # No additional loss parameters
    regularizer=None,  # No regularizer
    optimizer="Adam",  # Use the Adam optimizer
    optimizer_kwargs=dict(lr=0.01),  # Set the learning rate
    evaluator="RankBasedEvaluator",  # Use the rank-based evaluator
    evaluator_kwargs=dict(filtered=True),  # Use filtered evaluation
    evaluation_kwargs=dict(batch_size=256),  # Batch size for evaluation
)
```

Gives me 
```
1507     )
   1508 
-> 1509     model_instance = _handle_model(
   1510         device=device,
   1511         _result_tracker=_result_tracker,

[/usr/local/lib/python3.9/dist-packages/pykeen/pipeline/api.py](https://localhost:8080/#) in _handle_model(device, _result_tracker, _random_seed, training, model, model_kwargs, interaction, interaction_kwargs, dimensions, loss, loss_kwargs, regularizer, regularizer_kwargs)
    977         # TODO should kwargs for loss and regularizer be checked and raised for?
    978     else:
--> 979         model_instance, model_kwargs = _build_model_helper(
    980             model=model,
    981             model_kwargs=model_kwargs,

[/usr/local/lib/python3.9/dist-packages/pykeen/pipeline/api.py](https://localhost:8080/#) in _build_model_helper(model, model_kwargs, loss, loss_kwargs, _device, _random_seed, regularizer, regularizer_kwargs, training_triples_factory)
    844 
    845     return (
--> 846         model_resolver.make(
    847             model,
    848             triples_factory=training_triples_factory,

[/usr/local/lib/python3.9/dist-packages/class_resolver/api.py](https://localhost:8080/#) in make(self, query, pos_kwargs, **kwargs)
    202             cls: Type[X] = self.lookup(query)
    203             try:
--> 204                 return cls(**(pos_kwargs or {}), **kwargs)  # type: ignore
    205             except TypeError as e:
    206                 if "required keyword-only argument" in e.args[0]:

[/usr/local/lib/python3.9/dist-packages/pykeen/models/unimodal/compgcn.py](https://localhost:8080/#) in __init__(self, triples_factory, embedding_dim, encoder_kwargs, interaction, interaction_kwargs, **kwargs)
     69 
     70         # combined representation
---> 71         entity_representations, relation_representations = CombinedCompGCNRepresentations(
     72             triples_factory=triples_factory,
     73             **encoder_kwargs,

[/usr/local/lib/python3.9/dist-packages/pykeen/nn/representation.py](https://localhost:8080/#) in __init__(self, triples_factory, entity_representations, entity_representations_kwargs, relation_representations, relation_representations_kwargs, num_layers, dims, layer_kwargs)
    819         super().__init__()
    820         # TODO: Check
--> 821         assert triples_factory.create_inverse_triples
    822         self.entity_representations = build_representation(
    823             max_id=triples_factory.num_entities,

AssertionError:
```
If I replace OGBBioKG with "hetionet" then it is working fine. I wonder if this is another thing to check. Thanks a lot!

CompGCN needs to have inverse relations enabled. You can do this by passing `create_inverse_triples` to the dataset.

Thanks! It's all properly working now!