[OpenKE] Knowledge Embedding PyTorch版本

安装

1.  Install PyTorch
2.  Clone the OpenKE-PyTorch branch:
    $ git clone -b OpenKE-PyTorch [https://github.com/thunlp/OpenKE](https://github.com/thunlp/OpenKE)
    $ cd OpenKE
3.  Compile C++ files
    $ bash make.sh

训练

需要三个文件

• train2id.txt 训练文件，第一行是三元组的数量，接下来的数据格式是(e1,e2,rel),需要注意的是，e1,和e2是实体的编号，rel是关系的编号，其对应关系存放在文件entity2id.txt和relation2id.txt。
• entity2id.txt 第一行是实体的数目，接下来的每一行是实体对应相关的id。
• relation2id.txt 第一行是关系的数目，接下来的每一行是关系对应的id。

测试

需要5个文件，除过上面的三个，还需要

• test2id.txt 第一行是测试三元组的数目，接下来的是(e1,e2,rel).
• valid2id.txt 验证数据集，第一行是验证三元组的个数，接下来的行是验证数据(e1,e2,rel)
• type_constrain.txt 类型限制文件，表示关系只能和特定类型的头实体和尾实体结合，文件第一行是关系的数目，接下来的行是每一种关系的类型限制，例如： 某个relation的id是1200 ，他的头实体（head entities）有四种类型 3123 1034,58，5733 这个relation同时又有4种类型的尾实体12123,4388,11087,11088 ，这种n对n的关系可以通过through n-n.py in folder benchmarks/FB15K来查看。

Quick Start

import config
import models
import json
import numpy as np


con = config.Config()
#Input training files from benchmarks/FB15K/ folder.
con.set_in_path("./benchmarks/FB15K/")
con.set_work_threads(4)
con.set_train_times(500)
con.set_nbatches(100)
con.set_alpha(0.001)
con.set_margin(1.0)
con.set_bern(0)
con.set_dimension(50)
con.set_ent_neg_rate(1)
con.set_rel_neg_rate(0)
con.set_opt_method("SGD")

#Models will be exported via tf.Saver() automatically.
con.set_export_files("./res/model.vec.tf", 0)
#Model parameters will be exported to json files automatically.
con.set_out_files("./res/embedding.vec.json")
#Initialize experimental settings.
con.init()
#Set the knowledge embedding model
con.set_model(models.TransE)
#Train the model.
con.run()   

步骤1 加载数据

这个文件夹下面有三个文件train2id.txt,entity2id.txt,relation2id.txt。

con.set_in_path("benchmarks/FB15K/")

可以分配几个threads进行采样sample positive and negative cases.

con.set_work_threads(8)

步骤2, 设置训练参数

最大训练轮数，batchSeize，实体和关系的维数，

con.set_train_times(500)
con.set_nbatches(100)
con.set_alpha(0.5)
con.set_dimension(200)
con.set_margin(1)

对于负采样，我们可以把正常实体和关系拆分来构造negative triples， set_bern(0)是传统的采样方法，set_bern(1)表示使用 (Wang et al. 2014) denoted as "bern"提出的构造方法，set_ent_neg_rate是设置实体的负采样率，set_rel_neg_rate是设置关系的负采样率。

con.set_bern(0)
con.set_ent_neg_rate(1)
con.set_rel_neg_rate(0)

设置优化方法

con.set_optimizer("SGD")

步骤3，输出结果

模型参数每隔几轮就会使用torch.save()自动的保存下来，同时最终的结果会保存成json 文件的形式。

con.set_export_files("./res/model.vec.pt")
con.set_out_files("./res/embedding.vec.json")

步骤4 训练模型

con.init()
con.set_model(models.TransE)
con.run()

步骤5 测试

测试任务

link prediction任务：用于预测三元组中缺失的关系或者尾实体，对于测试的三元组，我们replace掉了head/tail 实体，并以降序的顺序给出预测出实体的得分。平均的指标有：

• MR：正确实体的平均rank。
• MRR: the average of the reciprocal ranks of correct entities。
• Hit@N：正确实体在top-N的比率

三元组分类任务：判断一个三元组(h,r,t)是否正确，这是一个二分类问题。
预测头实体任务： 预测topk个可能的头实体，所有的头实体用id表示

def predict_head_entity(self, t, r, k):
    r'''This mothod predicts the top k head entities given tail entity and relation.
    
    Args: 
        t (int): tail entity id
        r (int): relation id
        k (int): top k head entities
    
    Returns:
        list: k possible head entity ids        
    '''
    self.init_link_prediction()
    if self.importName != None:
        self.restore_pytorch()
    test_h = np.array(range(self.entTotal))
    test_r = np.array([r] * self.entTotal)
    test_t = np.array([t] * self.entTotal)
    res = self.trainModel.predict(test_h, test_t, test_r).data.numpy().reshape(-1).argsort()[:k]
    print(res)
    return res

预测尾实体：与预测头实体相似。
预测关系：

def predict_relation(self, h, t, k):
    r'''This methods predict the relation id given head entity and tail entity.
    
    Args:
        h (int): head entity id
        t (int): tail entity id
        k (int): top k relations
    
    Returns:
        list: k possible relation ids
    '''
    self.init_link_prediction()
    if self.importName != None:
        self.restore_pytorch()
    test_h = np.array([h] * self.relTotal)
    test_r = np.array(range(self.relTotal))
    test_t = np.array([t] * self.relTotal)
    res = self.trainModel.predict(test_h, test_t, test_r).data.numpy().reshape(-1).argsort()[:k]
    print(res)
    return res

预测三元组：给一个三元组，这个函数告诉我们是否这个三元组是否正确，如果threshold没有给出，那么函数从验证集中计算出这个关系的threshold。

def predict_triple(self, h, t, r, thresh = None):
    r'''This method tells you whether the given triple (h, t, r) is correct of wrong

    Args:
        h (int): head entity id
        t (int): tail entity id
        r (int): relation id
        thresh (fload): threshold for the triple
    '''
    self.init_triple_classification()
    if self.importName != None:
        self.restore_pytorch()  
    res = self.trainModel.predict(np.array([h]), np.array([t]), np.array([r])).data.numpy()
    if thresh != None:
        if res < thresh:
                        print("triple (%d,%d,%d) is correct" % (h, t, r))
                else:
                        print("triple (%d,%d,%d) is wrong" % (h, t, r)) 
        return
    self.lib.getValidBatch(self.valid_pos_h_addr, self.valid_pos_t_addr, self.valid_pos_r_addr, self.valid_neg_h_addr, self.valid_neg_t_addr, self.valid_neg_r_addr)
    res_pos = self.trainModel.predict(self.valid_pos_h, self.valid_pos_t, self.valid_pos_r)
    res_neg = self.trainModel.predict(self.valid_neg_h, self.valid_neg_t, self.valid_neg_r)
    self.lib.getBestThreshold(self.relThresh_addr, res_pos.data.numpy().__array_interface__['data'][0], res_neg.data.numpy().__array_interface__['data'][0])
    if res < self.relThresh[r]:
        print("triple (%d,%d,%d) is correct" % (h, t, r))
    else: 
        print("triple (%d,%d,%d) is wrong" % (h, t, r))

具体实施

第一步是导入数据集，然后配置参数，然后设置模型参数和测试的模型，例如我们需要测试TransE:有三种方法来测试模型。

1. 设置导入文件，OpenKE-PyTorch会自动的通过torch.load()来加载模型。

import config
import models
import numpy as np
import json

con = config.Config()
con.set_in_path("./benchmarks/FB15K/")
con.test_link_prediction(True)
con.test_triple_classification(True)
con.set_work_threads(4)
con.set_dimension(100)
con.set_import_files("./res/model.vec.pt")
con.init()
con.set_model(models.TransE)
con.test()

2. 从json文件中读取模型参数,手动的加载参数。

import config
import models
import numpy as np
import json

con = config.Config()
con.set_in_path("./benchmarks/FB15K/")
con.test_link_prediction(True)
con.test_triple_classification(True)
con.set_work_threads(4)
con.set_dimension(100)
con.init()
con.set_model(models.TransE)
f = open("./res/embedding.vec.json", "r")
content = json.loads(f.read())
f.close()
con.set_parameters(content)
con.test()

3. 使用torch.load()手动的加载模型。

import config
import models
import numpy as np
import json

con = config.Config()
con.set_in_path("./benchmarks/FB15K/")
con.test_link_prediction(True)
con.test_triple_classification(True)
con.set_work_threads(4)
con.set_dimension(100)
con.init()
con.set_model(models.TransE)
con.import_variables("./res/model.vec.pt")
con.test()

获取embedding 矩阵

有四种方式来获取embedding矩阵

1. 设置import 文件那么OpenKE 会自动的使用torch.load()加载模型.
   使用con.get_parameters()函数来获得list类型的embedding矩阵，可以通过con.get_parameters("numpy")获得numpy类型的参数。

import json
import numpy as py
import config
import models
con = config.Config()
con.set_in_path("./benchmarks/FB15K/")
con.test_link_prediction(True)
con.test_triple_classification(True)
con.set_work_threads(4)
con.set_dimension(100)
con.set_import_files("./res/model.vec.pt")
con.init()
con.set_model(models.TransE)
# Get the embeddings (numpy.array)
embeddings = con.get_parameters("numpy")
# Get the embeddings (python list)
embeddings = con.get_parameters()

2. 从json文件中获取

import json
import numpy as py
import config
import models
con = config.Config()
con.set_in_path("./benchmarks/FB15K/")
con.test_link_prediction(True)
con.test_triple_classification(True)
con.set_work_threads(4)
con.set_dimension(100)
con.init()
con.set_model(models.TransE)
f = open("./res/embedding.vec.json", "r")
embeddings = json.loads(f.read())
f.close()

3. 手动的使用torch.load()加载模型，但是获取embedding的方式还是一样的。

con = config.Config()
con.set_in_path("./benchmarks/FB15K/")
con.test_link_prediction(True)
con.test_triple_classification(True)
con.set_work_threads(4)
con.set_dimension(100)
con.init()
con.set_model(models.TransE)
con.import_variables("./res/model.vec.pt")
# Get the embeddings (numpy.array)
embeddings = con.get_parameters("numpy")
# Get the embeddings (python list)
embeddings = con.get_parameters()

4. 从一个训练好的模型中立即拿到embedding。

#Models will be exported via tf.Saver() automatically.
con.set_export_files("./res/model.vec.pt")
#Model parameters will be exported to json files automatically.
con.set_out_files("./res/embedding.vec.json")
#Initialize experimental settings.
con.init()
#Set the knowledge embedding model
con.set_model(models.TransE)
#Train the model.
con.run()
#Get the embeddings (numpy.array)
embeddings = con.get_parameters("numpy")
#Get the embeddings (python list)
embeddings = con.get_parameters()

接口

Config的接口

• def set_alpha(alpha = 0.001) 设置学习率
• def set_lmbda(lmbda = 0.0)，设置正则化前面的系数。

To set the degree of the regularization on the parameters

• set_train_times(self, times) 相当于设置epoch
• def sampling() 从正样本和负样本中采样一个batch
• def set_in_path(self, path) 读取数据
• def set_out_files(self, path) 当训练结束的时候将模型的参数变成json文件存储下来。
• def set_import_files(self, path) 模型所有的参数都可以用这个文件夹里面恢复。
• def set_export_steps(self, steps) 每隔多少步存储一次文件
• def save_pytorch(self) 使用torch.save来保存模型
• def import_variables(self, path = None) 恢复 tensorflow模型，相当于restore_tensorflow()
• def set_parameters(self, lists) 从jaon 文件中加载parameters
• def get_parameters(self, mode = "numpy") 获取模型参数也就是每个实体的embedding
• def set_model(model) 表示使用什么model进行knowledge embedding
• def set_log_on(flag = 1) 如果设置为1那么表示会打印loss函数

The framework will print loss values during training if flag = 1

class Config(object):
        
    #To set the learning rate
    def set_alpha(alpha = 0.001)
    
    #To set the degree of the regularization on the parameters
    def set_lmbda(lmbda = 0.0)
    
    #To set the gradient descent optimization algorithm (SGD, Adagrad, Adadelta, Adam)
    def set_optimizer(optimizer = "SGD")
    
    #To set the data traversing rounds
    def set_train_times(self, times)
    
    #To split the training triples into several batches, nbatches is the number of batches
    def set_nbatches(nbatches = 100)
    
    #To set the margin for the loss function
    def set_margin(margin = 1.0)
    
    #To set the dimensions of the entities and relations at the same time
    def set_dimension(dim)
    
    #To set the dimensions of the entities
    def set_ent_dimension(self, dim)
    
    #To set the dimensions of the relations
    def set_rel_dimension(self, dim)
    
    #To allocate threads for each batch sampling
    def set_work_threads(threads = 1)
    
    #To set negative sampling algorithms, unif (bern = 0) or bern (bern = 1)
    def set_bern(bern = 1)
    
    #For each positive triple, we construct rate negative triples by corrupt the entity
    def set_ent_neg_rate(rate = 1)
    
    #For each positive triple, we construct rate negative triples by corrupt the relation
    def set_rel_neg_rate(rate = 0)
    
    #To sample a batch of training triples, including positive and negative ones.
    def sampling()

    #To import dataset from the benchmark folder
    def set_in_path(self, path)
    
    #To export model parameters to json files when training completed
    def set_out_files(self, path)
    
    #To set the import files, all parameters can be restored from the import files
    def set_import_files(self, path)
    
    #To set the export file of model paramters, and export results every few rounds
    def set_export_files(self, path, steps = 0)

    #To export results every few rounds
    def set_export_steps(self, steps)

    #To save model via torch.save()
    def save_pytorch(self)

    #To restore model via torch.load()
    def restore_pytorch(self)

    #To export model paramters, when path is none, equivalent to save_tensorflow()
    def export_variables(self, path = None)

    #To import model paramters, when path is none, equivalent to restore_tensorflow()
    def import_variables(self, path = None)
    
    #To export model paramters to designated path
    def save_parameters(self, path = None)

    #To manually load parameters which are read from json files
    def set_parameters(self, lists)
    
    #To get model paramters, if using mode "numpy", you can get np.array , else you can get python lists
    def get_parameters(self, mode = "numpy")

    #To set the knowledge embedding model
    def set_model(model)
    
    #The framework will print loss values during training if flag = 1
    def set_log_on(flag = 1)

    #This is essential when testing
    def test_link_prediction(True)
    def test_triple_classification(True)

模型的接口

class Model(object)

    # in_batch = True, return [positive_head, positive_tail, positive_relation]
    # The shape of positive_head is [batch_size, 1]
    # in_batch = False, return [positive_head, positive_tail, positive_relation]
    # The shape of positive_head is [batch_size]
    get_positive_instance(in_batch = True)
    
    # in_batch = True, return [negative_head, negative_tail, negative_relation]
    # The shape of positive_head is [batch_size, negative_ent_rate + negative_rel_rate]
    # in_batch = False, return [negative_head, negative_tail, negative_relation]
    # The shape of positive_head is [(negative_ent_rate + negative_rel_rate) * batch_size]      
    get_negative_instance(in_batch = True)

    # in_batch = True, return all training instances with the shape [batch_size, (1 + negative_ent_rate + negative_rel_rate)]
    # in_batch = False, return all training instances with the shape [(negative_ent_rate + negative_rel_rate + 1) * batch_size]
    def get_all_instance(in_batch = False)

    # in_batch = True, return all training labels with the shape [batch_size, (1 + negative_ent_rate + negative_rel_rate)]
    # in_batch = False, return all training labels with the shape [(negative_ent_rate + negative_rel_rate + 1) * batch_size]
    # The positive triples are labeled as 1, and the negative triples are labeled as -1
    def get_all_labels(in_batch = False)
    
    #To calulate the loss
    def forward(self)

    # To define loss functions for knowledge embedding models
    def loss_func()
    
    # To define the prediction functions for knowledge embedding models
    def predict(self)

    def __init__(config)

#The implementation for TransE
class TransE(Model)

#The implementation for TransH  
class TransH(Model)

#The implementation for TransR
class TransR(Model)

#The implementation for TransD
class TransD(Model)

#The implementation for RESCAL
class RESCAL(Model)

#The implementation for DistMult
class DistMult(Model)                   

#The implementation for ComplEx
class ComplEx(Model)