Knowledge graphs (KGs) are data structures that store information about different entities (nodes) and their relations (edges). A common approach of using KGs in various machine learning tasks is to compute knowledge graph embeddings. DGL-KE is a high performance, easy-to-use, and scalable package for learning large-scale knowledge graph embeddings. The package is implemented on top of Deep Graph Library (DGL) and runs on CPU machine, GPU machine, as well as clusters with a set of popular models, including TransE, TransR, RESCAL, DistMult, ComplEx, and RotatE.
Figure: DGL-KE Overall Architecture
To install the latest version of DGL-KE run:
pip install dgl
pip install dglke
Train a transE model on FB15k dataset by running the following command:
dglke_train --model_name TransE_l2 --dataset FB15k --batch_size 1000 --neg_sample_size 200 --hidden_dim 400 \
--gamma 19.9 --lr 0.25 --max_step 3000 --log_interval 100 --batch_size_eval 16 --test -adv \
--regularization_coef 1.00E-09 --num_thread 1 --num_proc 8
This command will download the FB15k dataset, train the transE model on that, and save the trained embeddings into the file.
DGL-KE is designed for learning at scale. It introduces various novel optimizations that accelerate training on knowledge graphs with millions of nodes and billions of edges. Benchmarks on knowledge graphs with over 86M nodes and 338M edges shows that DGL-KE can compute embeddings in 100 minutes on an EC2 instance with 8 GPUs and 30 minutes on an EC2 cluster with 4 machines (48 cores/machine). These results represent a 2×-5× speedup over the best competing approaches.
Figure: DGL-KE vs GraphVite on FB15k
Figure: DGL-KE vs Pytorch-BigGraph on Freebase
Learn more details with our documentation!
This project is licensed under the Apache-2.0 License.