所用的数据集为FB15K和WN18，使用论文中的数据。论文： Translating Embeddings for Modeling Multi-relational Data Knowledge Graph Embedding by Translating on Hyperplanes Learning Entity and Relation Embeddings for Knowledge Graph Completion Knowledge Graph Embedding via Dynamic Mapping Matrix

Paperwithcode： https://github.com/thunlp/TensorFlow-TransX

数据集详情：

数据集	关系数	实体数	训练集样本数	验证集样本数	测试集样本数
WN18	18	40,9443	141,442	5,000	5,000
FB15K	1,345	14,951	483,142	50,000	59,071

注：论文中讲数据进行1-1、1-n等拆分出多个数据，实际仅用 train2id.txt, valid2id.txt,test2id.txt三个数据集。所给数据如下，点击查看数据来源，第一行为文件行数，在使用时应删掉第一行。

59071
453 1347 37
3136 4357 588
8663 4522 307
2404 8386 186
722 806 37
1248 10937 26
9182 1043 20

数据下载地址下载后解压,data/下包含整理后的FB15K和WN18数据。

环境要求

开发语言：Python3 使用框架：Mindspore requirements.txt

mindspore>=1.6.0
numpy
easydict
tqdm

快速入门

将代码和数据下载解压后，代码所在根目录记为<root_dir>。

基本配置

在训练或测试时，应指定配置信息，修改<root_dir>/train.py中的3项配置：

(1) 配置数据集文件夹data_dir，文件夹下应包含 train2id.txt,valid2id.txt,test2id.txt等三个文件，分别对应训练集、验证集和测试集。 (2) 数据集 dataset，可选 FB15K或WN18。

(3) 模型model，可选 transD, transE, transH或transR。

训练与测试：打开代码所在根目录

训练命令：python train.py
测试命令：python eval.py

说明1：训练和推理支持CPU/GPU/Ascend平台。

脚本说明

提供实现的细节

脚本和样例代码

目录结构如图所示：

各文件如下：

eval.py: 验证脚本
README.md: 本文档
requirements.txt: Python库配置
images/: README.md图等保存位置
src/config.py: 训练集和测试集配置
src/dataset.py: 数据集读取类
src/transD,E,H,R: 定义的知识图谱模型
val/testTransD,E,H,R: 定义的测试各个模型的对应类

脚本参数

由于训练集不同，其实体数、关系数、训练参数等均不同，因此将配置根据数据集划分为 fb15k_config 和 wn18_config，两者参数含义一致，此处以 fb15k_config 进行说明。每个数据集中再配置每个模型的参数。

fb15k_config = ed({
    # 全局参数
    "model": "transE",              # 模型，可选 {"transD", "transE", "transH", "transR"}
    "n_entity": 14951,              # 实体数量
    "n_relation": 1345,             # 关系数量
    
    # 训练参数
    "pre_model_path": "",           # 预训练模型保存路径
    "pretrained": False,            # 是否使用预训练模型，为True时需要pre_model_path指定正确路径
    "model_save_dir": "checkpoints/",   # 模型保存路径
    "log_save_file": "log.out",     # 训练日志保存路径
    "iterations": 1000000,          # 最大迭代次数
    "batch_size": 4096,             # 批处理大小
    "learning_rate": 0.001,         # 学习率
    
    # 模型参数
    ## transD参数
    "transD": {
        "n_entity_dim": 100,        # 实体编码维度
        "n_relation_dim": 100,      # 关系编码维度
        "margin": 1.0,              # 算法中计算损失时参数
        "norm": 1                   # 计算损失所用范数，可选{1,2}
    },
    ## transE参数
    "transE": {     
        "n_dim": 100,               # 编码维度，实体编码维度==关系编码维度
        "margin": 1.0,              # 算法中计算损失时参数
        "norm": 1                   # 计算损失所用范数，可选1或2
    },
    ## transH参数
    "transH": {
        "n_dim": 100,               # 编码维度，实体编码维度==关系编码维度
        "margin": 1.0,              # 算法中计算损失时参数
        "norm": 1                   # 计算损失所用范数
    },
    ## transR参数
    "transR": {
        "n_entity_dim": 100,        # 实体编码维度
        "n_relation_dim": 100,      # 关系编码维度
        "margin": 1.0,              # 算法中计算损失时参数
        "norm": 1                   # 计算损失所用范数
    }
})

训练过程

目前训练过程和评估只支持CPU/GPU，使用Ascend训练损失会不断上升，无法完成训练。

训练时，每次随机抽取 batch_size 个三元组，并进行corrupt，每次训练为一次iteration，最多训练iterations次，iterations训练参数可从 src/config.py 进行配置。

训练

训练方法：

打开 train.py 所在目录，运行：

python train.py

训练过程日志

以GPU在数据集WN18上训练transE为例，其部分日志如下：

# python train.py 
dataset = WN18, model = transR
read file path: /dataset/data/WN18/train2id.txt
data_loader initialized!
<=================== 开始训练 ===================>
epoch [1000], loss = 1331.6389914550782, 1000 iterations spend 0.12995397647221882 minutes!
epoch [2000], loss = 975.5102862548828, 1000 iterations spend 0.11221280495325724 minutes!
...
epoch [808000], loss = 1.239701193332672, 1000 iterations spend 0.10795061588287354 minutes!
epoch [809000], loss = 1.3451107335090637, 1000
...
loss_record:
 [1331.6389914550782, ..., 1.239701193332672]
>=================== 训练结束 ===================<

说明：其他数据集或模型的训练日志可从此处下载。

评估

评估过程

评估过程分为两步，首先指定评估数据集、评估模型等参数。具体参数如下：

  # 数据集目录
  data_dir = "/dataset/data/WN18/"    
  # 数据集:可选{"FB15K", "WN18"}     
  dataset = "WN18" 
  # 模型:可选{"transD", "transE", "transH", "transR"}
  model="transH"     
  # 训练好的参数，即实体等相关编码，为.ckpt文件
  param_path = "/model/fb15k_transE.ckpt" # 填入训练好的权重所在路径

评估命令

python eval.py

点击下载训练好的ckpt模型参数，模型文件名包含数据集、模型名、无filter的hits@10和mean-rank值等信息。

评估结果

对所给数据集，推理结果如下：

针对hits@10指标

模型	FB15K hits@10(%)(no filter)	WN18 hits@10(%)(no filter)	FB15K hits@10(%)(with filter)	WB18 hits@10(%)(with filter)
transD	44.82	74.17	58.97	86.19
transE	45.09	74.26	59.22	86.38
transH	45.36	78.74	59.54	91.65
transR	44.92	51.64	58.02	59.41

针对Mean-Rank指标

对于Mean-Rank指标，与filter无关

模型	FB15K mean-rank	WN18 mean-rank
transD	212.6	682.6
transE	208.6	676.4
transH	210.0	639.8
transR	270.1	606.5

运行结束后，会输出类似下面的内容：

# python eval.py
read file path: /dataset/data/FB15K/test2id.txt
100%|████████████████████████████████████████████████████████████████████| 1000/1000 [00:57<00:00, 17.29it/s]
**************************************************
model: transR, dataset: FB15K, test, is_filter: False
hits10 = 0.33, mean_rank = [491.052]
**************************************************

性能

训练性能

提供您训练性能的详细描述，例如finishing loss, throughput, checkpoint size等

对模型transD、transE、transH、transR使用CPU/GPU平台训练，参数如下。多个模型采用单核2U8G CPU或PCIE V100-32G GPU等训练。

参数	值
Resource	1*Ascend-910(32GB)	ARM: 24 vCPUs 96GB
AI Engine
Ascend-Powered-Engine mindspore_1.5.1-cann_5.0.2-py_3.7-euler_2.8.3-aarch64
Dataset	FB15K或WN18
Training Parameters	batch_size=4096(for FB15K), batch_size=1024(for WN18), iterations=$10^5至2\times 10^5$, learning_rate=0.001
Loss	L1 norm
Uploaded Date	09/14/2022 (month/day/year)
Total time	≈15h
Scripts	link

每次迭代（1个batch_size对应的样本）耗时：

模型	耗时-FB15K (min)	耗时-WN18 (min)	收敛后平均损失-FB15K	收敛后平均损失-WN18
transD	4.3	5.8	18.0	1.7
transE	1.2	2.3	17.2	1.6
transH	9.2	9.0	17.0	1.5
transR	13.9	8.3	21.2	1.1

说明：模型的实体和关系编码维度和后文的“推理性能”中一致。

推理性能

提供推理性能的详细描述，包括耗时，精度等

你可以参照如下模板

Parameters	Ascend
Model Version	transD, transE, transH, transR
Resource	1*Ascend-910(32GB), ARM: 24 vCPUs 96GB
Uploaded Date	09/20/2022 (month/day/year)
MindSpore Version	mindspore_1.5.1-py_3.7
Dataset	FB15K/WN18

推理编码维度与模型大小

模型	实体/关系编码维度（FB15K）	实体/关系编码维度（WN18）	模型大小（FB15K）	模型大小（WN18）	推理耗时-FB15K（ms/样本）	推理耗时-WN18（ms/样本）
transD	100 / 100	50 / 50	12.4M	15.6M	479	125
transE	100 / 100	50 / 50	6.21M	7.81M	452	85.5
transH	100 / 100	50 / 50	6.72M	7.81M	459	92.0
transR	100 / 100	50 / 50	57.5M	7.98M	461	85.7

说明1：表中推理耗时为无filter时测得结果，有filter时耗时为无filter的1.2倍左右。
说明2：总推理耗时与样本数量相关，表中显示单个样本平均耗时。FB15K测试集样本数为59071，WN18数据集样本数为5000，可据此估计完成所有样本耗时分别为6小时左右和10分钟左右。

贡献指南

如果你想参与贡献昇思的工作当中，请阅读昇思贡献指南和how_to_contribute

贡献者

此部分根据自己的情况进行更改，填写自己的院校和邮箱

email: forcekeng@126.com (同济大学)
gitee

ModelZoo 主页

请浏览官方主页。

forcekeng/TransX

目录

模型名称

模型架构

数据集