DGL-KE: Training Knowledge Graph Embeddings on CPU, GPU or a Cluster
High performance, easy-to-use, and scalable package for learning large-scale knowledge graph embeddings.
At a glance
- What is it?
- DGL-KE is an AWS Labs package built on Deep Graph Library that trains TransE, TransR, RESCAL, DistMult, ComplEx and RotatE embeddings, with a distributed trainer for graphs too large for one machine. The catch is maintenance: the last tagged release is 0.1.1 from August 2020, and the README now points new TransE, DistMult and RotatE users at GraphStorm instead.
- Who is it for?
- Adopt DGL-KE if you already have DGL in your stack, need one of the six listed models on a graph that fits a single machine, and can pin dgl and dglke to versions you have actually run together. Do not adopt it if you are starting fresh with TransE, DistMult or RotatE, because the README itself sends those users to GraphStorm, or if you need a model outside the listed six.
- Can I use it commercially?
- Yes. Apache-2.0 is a permissive licence: you can use, modify and sell software built on it, as long as you keep its copyright and licence notices.
- Is it still maintained?
- Yes. The repository last received commits 71 days ago.
- What is it written in?
- Mainly Python, according to GitHub's language statistics.
Answers come from the project's GitHub data, last synced on September 15, 2026, and from our analysis. They are not legal advice.
DEEP OPEN-SOURCE ANALYSIS
The problem DGL-KE addresses: embeddings for graphs that do not fit in memory
A knowledge graph stores entities as nodes and relations as edges, and the usual way to feed one into a machine learning pipeline is to compute embeddings for both. The hard part is scale. The README describes DGL-KE as a package for learning large-scale knowledge graph embeddings, implemented on top of Deep Graph Library, and its stated benchmark is a knowledge graph with over 86M nodes and 338M edges. That is the target: graphs where a single-process trainer either runs out of memory or takes long enough that iteration stops being practical. DGL-KE ships two training entry points for exactly this reason. dglke_train runs on one machine, CPU or GPU, and dglke_dist_train runs across a cluster. The same codebase covers both, so a model that works on a workstation is not rewritten for the cluster. The intended user is an engineer or researcher who already knows which embedding model they want and needs the training loop to finish. It is not a graph database, not a query engine, and not a service. It reads a graph, writes embedding files, and provides evaluation and inference commands over those files.
Six models, three tasks, and the commands that map to them
The README lists TransE, TransR, RESCAL, DistMult, ComplEx and RotatE as the supported models, and three tasks. Training uses dglke_train on a single machine or dglke_dist_train in a distributed environment. Evaluation reads pre-trained embeddings and scores them with a link prediction task on the test set, via dglke_eval. Inference has two forms: dglke_predict for entity and relation linkage prediction, and dglke_emb_sim for embedding similarity. That split matters when you plan a deployment, because evaluation and inference both consume saved embedding files rather than re-running training. You can train once, keep the output, and run dglke_eval or dglke_predict against it as many times as you like. The model list is also a boundary. If your work needs a scoring function outside those six, this package does not offer a plug-in path that the README describes, so you would be modifying the training code rather than configuring it.
How the distributed trainer splits work across machines
The README does not document the internals of the distributed path beyond naming dglke_dist_train and pointing at the SIGIR 2020 paper for the optimizations. What the repository does make visible is the dependency structure: DGL-KE is implemented on top of DGL, and the quick start sets DGLBACKEND=pytorch before invoking training. So the graph sampling and the message-passing machinery come from DGL, and DGL-KE supplies the knowledge graph specific pieces: the negative sampling, the scoring functions for the six models, and the distributed training loop. The benchmark claim in the README is that DGL-KE computes embeddings for the 86M node, 338M edge graph in 100 minutes on an EC2 instance with 8 GPUs and 30 minutes on an EC2 cluster with 4 machines at 48 cores each, described as a 2x to 5x speedup over the best competing approaches. Those numbers come from the project's own benchmark and its paper; they are not independently reproduced here, and the hardware is specific enough that you should treat them as an upper bound on what a different cluster will do.
Getting it running: install, backend variable, and a real training command
Installation is two pip commands from the README: sudo pip3 install dgl followed by sudo pip3 install dglke. Note the ordering and the absence of version pins. DGL-KE is built on DGL, and nothing in the README states which DGL versions the 0.1.1 release was tested against, so the install you get today may pair a recent dgl with a 2020 dglke. The quick start then trains TransE on FB15k with a single command, prefixed by DGLBACKEND=pytorch to select the backend. The flags shown are --model_name TransE_l2, --dataset FB15k, --batch_size 1000, --neg_sample_size 200, --hidden_dim 400, --gamma 19.9, --lr 0.25, --max_step 500, --log_interval 100, --batch_size_eval 16, -adv, --regularization_coef 1.00E-09, --test, --num_thread 1 and --num_proc 8. The README states that this downloads the FB15k dataset, trains the model and saves the trained embeddings to a file. The dataset argument is doing double duty: it names a known graph and triggers the download. For your own data you would be supplying files rather than a dataset name, and the README excerpt here does not cover that format, so check the documentation site before assuming the flag behaves the same way.
The maintenance signal: a 2020 release tag and a redirect to GraphStorm
The most important line in the README is the note above the quick start. It says that if you just want to train a KGE model using TransE, DistMult or RotatE, you should go to awslabs/graphstorm. Read that carefully. Three of the six advertised models are the ones the maintainers now point elsewhere for. The last tagged releases are 0.1.1 from August 2020 and 0.1.0 from April 2020, while the repository shows pushes in 2026, so the code is not frozen but the release process is. In practice that means you install from PyPI at a version five years old, or you install from the master branch and accept whatever is on it. Neither is unusual for research infrastructure, but it changes how you plan upgrades: there is no changelog cadence to follow, and no stated deprecation timeline for the package. The Apache-2.0 licence is permissive and imposes no copyleft obligation on your own code, but it also comes with no warranty, and this is not legal advice; if you redistribute DGL-KE or a modified version, read the licence text and your organisation's policy rather than assuming the badge on the README settles it.
Where DGL-KE is the wrong tool, and what to use instead
DGL-KE is the wrong choice when your graph fits comfortably on one machine and you are not already using DGL. The dependency on DGL is the whole reason the distributed path exists, and if you do not need that path you are paying for it in install complexity and in coupling to DGL's release cycle. The README's own redirect makes the alternative concrete: GraphStorm, from the same awslabs organisation, is where the project sends users who want TransE, DistMult or RotatE. That is not a competitor in the usual sense; it is the successor path the maintainers chose, and the difference in approach is that GraphStorm is positioned as the place to go for those models rather than as a drop-in replacement for the full DGL-KE command set. The README also names two systems it benchmarked against, GraphVite on FB15k and Pytorch-BigGraph on Freebase, and shows comparison figures. If you are choosing between DGL-KE and Pytorch-BigGraph specifically, the meaningful difference is that Pytorch-BigGraph is a standalone trainer while DGL-KE inherits its graph handling from DGL, which matters if you already have DGL pipelines. The README does not give enough detail on either comparison to decide on anything but your own measurements.
Who should adopt DGL-KE, and what to check before you do
Adopt it if you have a graph in the tens of millions of nodes or more, you need one of the six listed models, and either a single multi-GPU machine or a cluster is available to you. The distributed trainer and the shared code path between single-machine and cluster runs are the reason to pick this over assembling a training loop yourself. Do not adopt it if you are starting a new project today around TransE, DistMult or RotatE, because the README routes those users to GraphStorm, or if you need a model outside the six. Before you commit, verify three things. First, that the dgl and dglke versions you install actually work together, since the quick start pins neither and the newest dglke tag is from 2020. Second, that FB15k downloads and completes a short run with the exact command above on your hardware, which is the cheapest end-to-end check available. Third, whether you are installing the 0.1.1 release or tracking master, because those are different codebases and the README does not describe what changed between them.
Editorial conclusion
Adopt DGL-KE if you already have DGL in your stack, need one of the six listed models on a graph that fits a single machine, and can pin dgl and dglke to versions you have actually run together. Do not adopt it if you are starting fresh with TransE, DistMult or RotatE, because the README itself sends those users to GraphStorm, or if you need a model outside the listed six. Before committing, verify that the dgl and dglke versions you install are compatible, that FB15k downloads and trains through dglke_train with your chosen backend, and whether the 0.1.1 tag is the commit you want given that the repository is still being pushed to.
Community notes