search

cai-lw / KBGAN

Code for "KBGAN: Adversarial Learning for Knowledge Graph Embeddings" https://arxiv.org/abs/1711.04071
MIT License
212 stars 57 forks source link
generative-adversarial-network knowledge-base knowledge-base-embeddings knowledge-graph knowledge-graph-embeddings reinforcement-learning

readme

KBGAN

Liwei Cai and William Yang Wang, "KBGAN: Adversarial Learning for Knowledge Graph Embeddings", in Proceedings of The 16th Annual Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (NAACL HLT 2018).

Paper: https://arxiv.org/abs/1711.04071

Our lab: http://nlp.cs.ucsb.edu/index.html

Dependencies

PyTorch 0.2.0 is REQUIRED. Because PyTorch is not backward compatible, newer versions will NOT work. We understand that 0.2.0 is outdated, but we currently have no schedule of adding support to newer PyTorch versions.

Usage

  1. Unzip data.zip.
  2. Pretrain: python3 pretrain.py --config=config_<dataset_name>.yaml --pretrain_config=<model_name> (this will generate a pretrained model file)
  3. Adversarial train: python3 gan_train.py --config=config_<dataset_name>.yaml --g_config=<G_model_name> --d_config=<D_model_name> (make sure that G model and D model are both pretrained)

Feel free to explore and modify parameters in config files. Default parameters are those used in experiments reported in the paper.

Decrease test_batch_size in config files if you experience GPU memory exhaustion. (this would make the program runs slower, but would not affect the test result)