[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:有三种方法来测试模型。

设置导入文件，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()

从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()

使用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矩阵

设置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()

从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()

手动的使用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()

从一个训练好的模型中立即拿到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)

最后编辑于：2018.12.02 18:42:38

人面猴
序言：七十年代末，一起剥皮案震惊了整个滨河市，随后出现的几起案子，更是在滨河造成了极大的恐慌，老刑警刘岩，带你破解...
沈念sama阅读 211,376评论 6赞 491
死咒
序言：滨河连续发生了三起死亡事件，死亡现场离奇诡异，居然都是意外死亡，警方通过查阅死者的电脑和手机，发现死者居然都...
沈念sama阅读 90,126评论 2赞 385
救了他两次的神仙让他今天三更去死
文/潘晓璐我一进店门，熙熙楼的掌柜王于贵愁眉苦脸地迎上来，“玉大人，你说我怎么就摊上这事。” “怎么了？”我有些...
开封第一讲书人阅读 156,966评论 0赞 347
道士缉凶录：失踪的卖姜人
文/不坏的土叔我叫张陵，是天一观的道长。经常有香客问我，道长，这世上最难降的妖魔是什么？我笑而不...
开封第一讲书人阅读 56,432评论 1赞 283
港岛之恋（遗憾婚礼）
正文为了忘掉前任，我火速办了婚礼，结果婚礼上，老公的妹妹穿的比我还像新娘。我一直安慰自己，他们只是感情好，可当我...
茶点故事阅读 65,519评论 6赞 385
恶毒庶女顶嫁案：这布局不是一般人想出来的
文/花漫我一把揭开白布。她就那样静静地躺着，像睡着了一般。火红的嫁衣衬着肌肤如雪。梳的纹丝不乱的头发上，一...
开封第一讲书人阅读 49,792评论 1赞 290
城市分裂传说
那天，我揣着相机与录音，去河边找鬼。笑死，一个胖子当着我的面吹牛，可吹牛的内容都是我干的。我是一名探鬼主播，决...
沈念sama阅读 38,933评论 3赞 406
双鸳鸯连环套：你想象不到人心有多黑
文/苍兰香墨我猛地睁开眼，长吁一口气：“原来是场噩梦啊……” “哼！你这毒妇竟也来了？” 一声冷哼从身侧响起，我...
开封第一讲书人阅读 37,701评论 0赞 266
万荣杀人案实录
序言：老挝万荣一对情侣失踪，失踪者是张志新（化名）和其女友刘颖，没想到半个月后，有当地人在树林里发现了一具尸体，经...
沈念sama阅读 44,143评论 1赞 303
护林员之死
正文独居荒郊野岭守林人离奇死亡，尸身上长有42处带血的脓包…… 初始之章·张勋以下内容为张勋视角年9月15日...
茶点故事阅读 36,488评论 2赞 327
白月光启示录
正文我和宋清朗相恋三年，在试婚纱的时候发现自己被绿了。大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
茶点故事阅读 38,626评论 1赞 340
活死人
序言：一个原本活蹦乱跳的男人离奇死亡，死状恐怖，灵堂内的尸体忽然破棺而出，到底是诈尸还是另有隐情，我是刑警宁泽，带...
沈念sama阅读 34,292评论 4赞 329
日本核电站爆炸内幕
正文年R本政府宣布，位于F岛的核电站，受9级特大地震影响，放射性物质发生泄漏。R本人自食恶果不足惜，却给世界环境...
茶点故事阅读 39,896评论 3赞 313
男人毒药：我在死后第九天来索命
文/蒙蒙一、第九天我趴在偏房一处隐蔽的房顶上张望。院中可真热闹，春花似锦、人声如沸。这庄子的主人今日做“春日...
开封第一讲书人阅读 30,742评论 0赞 21
一桩弑父案，背后竟有这般阴谋
文/苍兰香墨我抬头看了看天上的太阳。三九已至，却和暖如春，着一层夹袄步出监牢的瞬间，已是汗流浃背。一阵脚步声响...
开封第一讲书人阅读 31,977评论 1赞 265
情欲美人皮
我被黑心中介骗来泰国打工，没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留，地道东北人。一个月前我还...
沈念sama阅读 46,324评论 2赞 360
代替公主和亲
正文我出身青楼，却偏偏与公主长得像，于是被迫代替她去往敌国和亲。传闻我的和亲对象是个残疾皇子，可洞房花烛夜当晚...
茶点故事阅读 43,494评论 2赞 348