完整源码:zx.xcxyms.top
GraphVite提供了一个更快的RotatE和其他KGE模型的多gpu实现。
实现的功能
√模型:
√pRotatE
√TransE
√复杂的
√DistMult
评价指标:
√MRR, MR, HITS@1, HITS@3, HITS@10(过滤)
√AUC-PR(国家数据集)

损失函数:
√统一-抽样
√Self-Adversarial负采样
使用
知识付费源码的数据:
●entities.dict:映射实体到唯一id的字典
●relations.dict:映射到唯一id的字典
●train.txt:训练KGE模型拟合该数据集
●valid.txt:如果没有验证数据,则创建一个空白文件
●用法: KGE模型在此数据集上进行评估
例如,该命令在GPU 0的FB15k数据集上训练一个RotatE模型。
CUDA_VISIBLE_DEVICES=0 python -u codes/run.py --do_train \
--cuda \
--do_valid \
--do_test \
--data_path data/FB15k \
--model RotatE \
-n 256 -b 1024 -d 1000 \
-g 24.0 -a 1.0 -adv \
-lr 0.0001 --max_steps 150000 \
-save models/RotatE_FB15k_0 --test_batch_size 16 -de
查看code /run.py的argparse配置以获取更多参数和详细信息。
测试
CUDA_VISIBLE_DEVICES=$GPU_DEVICE python -u $CODE_PATH/run.py——do_test——cuda -init $SAVE
、为了再现ICLR 2019论文《RotatE: Knowledge Graph em
run.sh脚本提供了一种简单的方法来搜索超参数:
bash run.sh train RotatE FB15k 00 1024 256 1000 24.0 1.0 0.0001 200000 16 -de
速度
KGE模型通常需要大约半个小时在一个默认配置的GeForce GTX 1080 Ti GPU上运行10000步。这些模型需要不同的max_steps来收敛不同的数据集:
python库由3个对象组成:
●TrainDataset (dataloader.py):为培训准备数据流
●TestDataSet (dataloader.py):为评估准备数据流
●KGEModel (model.py):计算三重评分并提供训练/测试API
py文件包含main函数,它解析参数、读取数据、初始化模型并提供训练循环。
def TransE(self, head, relation, tail, mode):
if mode == 'head-batch':
score = head + (relation - tail)
else:
score = (head + relation) - tail
score = self.gamma.item() - torch.norm(score, p=1, dim=2)
return score
引用
如果您使用代码,请引用以下:
@inproceedings{
sun2018rotate,
ti
author={Zhiqing Sun and Zhi-Hong Deng and Jian-Yun Nie and Jian Tang},
bookti
year={2019},
url={/forum?id=HkgEQnRqYQ},
}
像这样添加你自己的模型到model.py:
最佳答案
