{"id":28700533,"url":"https://github.com/deepgraphlearning/ultra_torchdrug","last_synced_at":"2025-06-14T11:08:20.636Z","repository":{"id":203203416,"uuid":"709051885","full_name":"DeepGraphLearning/ultra_torchdrug","owner":"DeepGraphLearning","description":"A TorchDrug version of ULTRA for reproducibility","archived":false,"fork":false,"pushed_at":"2023-10-29T07:23:13.000Z","size":13105,"stargazers_count":5,"open_issues_count":2,"forks_count":2,"subscribers_count":4,"default_branch":"main","last_synced_at":"2025-01-06T09:03:26.356Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":null,"language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/DeepGraphLearning.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null}},"created_at":"2023-10-23T23:07:30.000Z","updated_at":"2024-11-13T19:13:04.000Z","dependencies_parsed_at":null,"dependency_job_id":"55df40d3-28fb-434e-9b88-6f3d76226584","html_url":"https://github.com/DeepGraphLearning/ultra_torchdrug","commit_stats":null,"previous_names":["deepgraphlearning/ultra_torchdrug"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/DeepGraphLearning/ultra_torchdrug","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/DeepGraphLearning%2Fultra_torchdrug","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/DeepGraphLearning%2Fultra_torchdrug/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/DeepGraphLearning%2Fultra_torchdrug/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/DeepGraphLearning%2Fultra_torchdrug/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/DeepGraphLearning","download_url":"https://codeload.github.com/DeepGraphLearning/ultra_torchdrug/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/DeepGraphLearning%2Fultra_torchdrug/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":259804865,"owners_count":22913903,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2025-06-14T11:08:18.296Z","updated_at":"2025-06-14T11:08:20.627Z","avatar_url":"https://github.com/DeepGraphLearning.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"\u003cdiv align=\"center\"\u003e\n\n# ULTRA: Towards Foundation Models for Knowledge Graph Reasoning #\n\n[![pytorch](https://img.shields.io/badge/PyTorch_1.13+-ee4c2c?logo=pytorch\u0026logoColor=white)](https://pytorch.org/get-started/locally/)\n[![arxiv](http://img.shields.io/badge/arxiv-2310.04562-yellow.svg)](https://arxiv.org/abs/2310.04562)\n![license](https://img.shields.io/badge/License-MIT-green.svg?labelColor=gray)\n\n\u003c/div\u003e\n\nThis is the TorchDrug implementation of [ULTRA], a foundation model for KG reasoning. Authored by [Michael Galkin], [Zhaocheng Zhu], and [Xinyu Yuan]. This repo contains the original code to reproduce the experimental results reported in the paper. The latest maintained version of ULTRA is available in the [PyG version of ULTRA](https://github.com/DeepGraphLearning/ULTRA).\n\n[Zhaocheng Zhu]: https://kiddozhu.github.io\n[Michael Galkin]: https://migalkin.github.io/\n[Xinyu Yuan]: https://github.com/KatarinaYuan\n[Ultra]: https://deepgraphlearning.github.io/project/ultra\n\n## Installation ##\n\nYou may install the dependencies via either conda or pip. \nUltra (TorchDrug) is compatible with Python 3.7/3.8/3.9, PyTorch 1.13 and PyG 2.3 (CUDA 11.7 or later wnen running on GPUs). If you are on a Mac, you may omit the CUDA toolkit requirements (tested with PyTorch 2.0 with the relevant `torch-scatter` version on Mac M2).\n\n### From Conda ###\n\n```bash\nconda install torchdrug pytorch cudatoolkit -c milagraph -c pytorch -c pyg\nconda install pytorch-sparse pytorch-scatter -c pyg\nconda install easydict pyyaml -c conda-forge\n```\n\n### From Pip ###\n\n```bash\npip install torchdrug torch\npip install easydict pyyaml\n```\n\nIf anything else is missing, install those from `requirements.txt`\n\n\u003cdetails\u003e\n\u003csummary\u003e Compilation of the `rspmm` kernel \u003c/summary\u003e\n\nTo make relational message passing iteration `O(V)` instead of `O(E)` we ship a custom `rspmm` kernel that will be compiled automatically upon the first launch. The `rspmm` kernel supports `transe` and `distmult` message functions, others like `rotate` will resort to full edge materialization and `O(E)` complexity.\n\nThe kernel can be compiled on both CPUs (including M1/M2 on Macs) and GPUs (it is done only once and then cached). For GPUs, you need a CUDA 11.7+ toolkit with the `nvcc` compiler. If you are deploying this in a Docker container, make sure to start from the `devel` images that contain `nvcc` in addition to plain CUDA runtime.\n\nMake sure your `CUDA_HOME` variable is set properly to avoid potential compilation errors, eg\n```bash\nexport CUDA_HOME=/usr/local/cuda-11.7/\n```\n\n\u003c/details\u003e\n\n\n## Checkpoints ##\n\nWe provide two pre-trained ULTRA checkpoints in the `/ckpts` folder of the same model size (6-layer GNNs per relation and entity graphs, 64d, 168k total parameters) trained on 2 x A100 GPUs with this codebase:\n* `td_ultra_3g.pth`: trained on `FB15k237, WN18RR, CoDExMedium` for 200,000 steps, config is in `/config/transductive/pretrain_3g.yaml`\n* `td_ultra_4g.pth`: trained on `FB15k237, WN18RR, CoDExMedium, NELL995` for 400,000 steps, config is in `/config/transductive/pretrain_4g.yaml`\n\nYou can use those checkpoints for zero-shot inference on any graph (including your own) or use it as a backbone for fine-tuning.\n\n## Run Inference and Fine-tuning\n\nThe `/scripts` folder contains 2 executable files:\n* `run_full.py` - run an experiment on a single dataset and/or pre-training;\n* `run_many.py` - run experiments on several datasets sequentially and dump results into a CSV file.\n\nThe yaml configs in the `config` folder are provided for both `transductive` and `inductive` datasets.\n\n### Run a single experiment\n\nThe `run_full.py` command requires the following arguments:\n* `-c \u003cyaml config\u003e`: a path to the yaml config\n* `--dataset`: dataset name (from the list of [datasets](#datasets))\n* `--version`: a version of the inductive dataset (see all in [datasets](#datasets)), not needed for transductive graphs. For example, `--dataset FB15k237Inductive --version v1` will load one of the GraIL inductive datasets.\n* `--epochs`: number of epochs to train, `--epochs 0` means running zero-shot inference.\n* `--bpe`: batches per epoch (replaces the length of the dataloader as default value). `--bpe 100 --epochs 10` means that each epoch consists of 100 batches, and overall training is 1000 batches. Set `--bpe null` to use the full length dataloader or comment the `bpe` line in the yaml configs.\n* `--gpus`: number of gpu devices, set to `--gpus null` when running on CPUs, `--gpus [0]` for a single GPU, or otherwise set the number of GPUs for a [distributed setup](#distributed-setup)\n* `--ckpt`: path to the one of the ULTRA checkpoints to use (you can use those provided in the repo ot trained on your own). Use `--ckpt null` to start training from scratch (or run zero-shot inference on a randomly initialized model, it still might surprise you and demonstrate non-zero performance).\n\nZero-shot inference setup is `--epochs 0` with a given checkpoint `ckpt`.\n\nFine-tuning of a checkpoint is when epochs \u003e 0 with a given checkpoint.\n\n\nAn example command for an inductive dataset to run on a CPU: \n\n```bash\npython script/run_full.py -c config/inductive/inference.yaml --dataset FB15k237Inductive --version v1 --epochs 0 --bpe null --gpus null --ckpt ckpts/ultra_4g.pth\n```\n\nAn example command for a transductive dataset to run on a GPU:\n```bash\npython script/run_full.py -c config/transductive/inference.yaml --dataset CoDExSmall --epochs 0 --bpe null --gpus [0] --ckpt ckpts/ultra_4g.pth\n```\n\n### Run on many datasets\n\nThe `run_many.py` script is a convenient way to run evaluation (0-shot inference and fine-tuning) on several datasets sequentially. Upon completion, the script will generate a csv file `ultra_results_\u003ctimestamp\u003e` with the test set results and chosen metrics. \nUsing the same config files, you only need to specify:\n\n* `-c \u003cyaml config\u003e`: use the full path to the yaml config because workdip will be reset after each dataset; \n* `-d, --datasets`: a comma-separated list of [datasets](#datasets) to run, inductive datasets use the `name:version` convention. For example, `-d FB15k237Inductive:v1,FB15k237Inductive:v2`;\n* `--ckpt`: ULTRA checkpoint to run the experiments on, use the full path to the file;\n* `--gpus`: the same as in [run single](#run-a-single-experiment);\n* `-reps` (optional): number of repeats with different seeds, set by default to 1 for zero-shot inference;\n* `-ft, --finetune` (optional): use the finetuning configs of ULTRA (`default_finetuning_config`) to fine-tune a given checkpoint for specified `epochs` and `bpe`;\n* `-tr, --train` (optional): train ULTRA from scratch on the target dataset taking `epochs` and `bpe` parameters from another pre-defined config (`default_train_config`);\n* `--epochs` and `--bpe` will be set according to a configuration, by default they are set for a 0-shot inference.\n\nAn example command to run 0-shot inference evaluation of an ULTRA checkpoint on 4 FB GraIL datasets:\n\n```bash\npython script/run_many.py -c /path/to/config/inductive/inference.yaml --gpus [0] --ckpt /path/to/ultra/ckpts/ultra_4g.pth -d FB15k237Inductive:v1,FB15k237Inductive:v2,FB15k237Inductive:v3,FB15k237Inductive:v4\n```\n\nAn example command to run fine-tuning on 4 FB GraIL datasets with 5 different seeds:\n\n```bash\npython script/run_many.py -c /path/to/config/inductive/inference.yaml --gpus [0] --ckpt /path/to/ultra/ckpts/ultra_4g.pth --finetune --reps 5 -d FB15k237Inductive:v1,FB15k237Inductive:v2,FB15k237Inductive:v3,FB15k237Inductive:v4\n```\n\n### Pretraining\n\nRun the script `run_full.py` with the `config/transductive/pretrain_3g.yaml` config file. \n\n`graphs` in the config specify the pre-training mixture. `pretrain_3g.yaml` uses FB15k237, WN18RR, CoDExMedium. By default, we use the training option `fast_test: 500` to run faster evaluation on a random subset of 500 triples (that approximates full validation performance) of each validation set of the pre-training mixture.\nYou can change the pre-training length by varying batches per epoch `batch_per_epoch` and `epochs` hyperparameters.\n\nAn example command to start pre-training on 3 graphs:\n\n```bash\npython script/run_full.py -c /path/to/config/transductive/pretrain_3g.yaml --gpus [0]\n```\n\nPre-training can be computationally heavy, you might need to decrease the batch size for smaller GPU RAM. The two provided checkpoints were trained on 4 x A100 (40 GB).\n\n#### Distributed setup\nTo run ULTRA with multiple GPUs, use the following commands (eg, 4 GPUs per node)\n\n```bash\npython -m torch.distributed.launch --nproc_per_node=4 script/run_full.py -c config/transductive/pretrain.yaml --gpus [0,1,2,3]\n```\n\nMulti-node setup might work as well(not tested):\n```bash\npython -m torch.distributed.launch --nnodes=4 --nproc_per_node=4 script/pretrain.py -c config/transductive/pretrain.yaml --gpus [0,1,2,3,0,1,2,3,0,1,2,3,0,1,2,3]\n```\n\n\n## Citation ##\n\nIf you find this codebase useful in your research, please cite the original paper.\n\n```bibtex\n@article{galkin2023ultra,\n  title={Towards Foundation Models for Knowledge Graph Reasoning},\n  author={Mikhail Galkin and Xinyu Yuan and Hesham Mostafa and Jian Tang and Zhaocheng Zhu},\n  year={2023},\n  eprint={2310.04562},\n  archivePrefix={arXiv},\n  primaryClass={cs.CL}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fdeepgraphlearning%2Fultra_torchdrug","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fdeepgraphlearning%2Fultra_torchdrug","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fdeepgraphlearning%2Fultra_torchdrug/lists"}