简单实现几篇知识图谱嵌入(Knowledge Graph Embedding，KGE)模型

关于知识图谱嵌入的理论介绍：
简要总结一篇关于知识图谱嵌入的综述

KGE的诸多方法

KGE就是将实体和关系嵌入到低维向量空间中，同时保留KG的结构和语义信息

现有的KGE方法可以划分为三类：

1. 基于翻译距离的(translational distance based)
2. 基于语义匹配的(semantic matching based)
3. 基于神经网络的(neural network based)

下面我们来实现几个最经典的KGE模型：基于翻译距离的模型TransE、基于语义匹配的模型RESCAL和DistMult，

TransE

该模型将关系看作头实体到尾实体的翻译。
TransE受到word2vec的启发，如果你已经训练好了词向量，那么针对三个单词：国家country、城市city，首都captial-of，就会有如下关系：$V_{country}-V_{city}$得到的向量与$V_{captial-of}$这个向量就很相近。TransE也认为，在KG的embedding空间中，也一定存在这种关系，即两个实体的嵌入向量的差值（或者其它操作）就代表这两个实体之间的关系。

因为关系和实体都被表示为向量，所以另一种数学化的说法就是在向量(vector)空间中，TransE将关系看作是头实体到尾实体的平移操作，即：
$$v_h+v_r\approx v_t$$

数学定义中：
平移前点的坐标+平移向量的坐标=平移后点的坐标
所以我们称之为平移

TransE的得分函数就是向量之间的欧氏距离的相反数：
$$Score(h,r,t)=-||v_h+v_r-v_t|{|}_2^2$$
损失函数定义为：
$$max(0,\gamma -Score(h,r,t)+Score(h^{\prime },r,t^{\prime }))$$
也就是在embedding space中，正例三元组的得分要比负例三元组的得分高出$\gamma$，又由于得分函数表示为距离的相反数，所以得分高代表距离近。即：正例三元组的距离要比负例三元组的距离小至少$\gamma$长度的距离。

RESCAL

该模型的得分函数定义为：
$$Score(h,r,t)=v_h^T{M}_r{v}_t$$

其中$v_h\in R^d,v_t\in R^d,M_r\in R^{d\times d}$
$v_h,v_t$都是从实体embedding矩阵（记为$E$）中的取出(根据实体id获取)的vector，$E$的形状是(num_entities,d)。
$M_r$是整个关系tensor（三维的）中的根据关系id获取的matrix（二维的）。整个关系tensor记为$R$，形状是(num_relations,d,d)，所以$M_r$的shape是($d,d$)。

DistMult

DistMult是RESCAL的简化。具体来说就是RESCAL中每一个head和tail实体之间的关系r是用一个matrix表示。而DistMult中则用一个vector表示两个实体间的关系。
所以得分函数是三个vector之间的内积：
$$<v_h,v_r,v_t>$$
$v_h,v_r,v_t\in R^d$

代码实现

获取数据

我们使用FB15K知识库

下载解压：

FB15K知识库就是TransE这篇论文的作者从Freebase知识库中选取的一部分三元组构成的一个小规模的知识库

三元组数量	实体数量	关系数量
592213	14951	1345

592213个三元组的划分情况是：

数据集划分	三元组数量
训练集	483142
验证集	50000
测试集	59017

实现

实现代码需要说明的一点是，下面的代码采用Binary Cross Entropy loss作为损失函数，即输入是头实体和关系的id，经过模型之后，输出一个向量，长度是num_entities，也就是所有实体的数量。

这个向量的每一个值都进行sigmoid运算，此时这个向量的第i个位置的值代表模型预测第i个实体是尾实体的概率

标签是尾实体的id，即告诉模型第几个位置的实体才是真正的尾实体，需要增加这个位置的概率，降低其余位置的概率。

也就是说达到了：要求正例三元组的得分大于负例三元组的得分。

导包

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import time
from collections import defaultdict
import argparse
from tqdm import tqdm
import os

处理数据

加载数据

def load_data(data_dir,data_type):
    with open("%s%s.txt" % (data_dir, data_type), "r") as f:
        data = f.read().strip().split("\n")
        data = [i.split('\t') for i in data]
        print(len(data),data_type)
        return data

train_data=load_data(data_dir='/mnt/cfs/speech/nlp/work/xhsun/KGQA/Trans/FB15k/',data_type='freebase_mtr100_mte100-train')
valid_data=load_data(data_dir='/mnt/cfs/speech/nlp/work/xhsun/KGQA/Trans/FB15k/',data_type='freebase_mtr100_mte100-valid')
test_data=load_data(data_dir='/mnt/cfs/speech/nlp/work/xhsun/KGQA/Trans/FB15k/',data_type='freebase_mtr100_mte100-test')
data=train_data+valid_data+test_data
print(len(data))

统计所有的头实体、尾实体以及 关系：

entities = sorted(list(set([d[0] for d in data]+[d[2] for d in data])))
print(len(entities))
relations = sorted(list(set([d[1] for d in data])))
print(len(relations))

即14951个实体，1345个关系

构造entity2id和relation2id的字典映射

entity_idxs={entities[i]:i for i in range(len(entities))}
relation_idxs={relations[i]:i for i in range(len(relations))}

构造实体到id，关系到id的映射，这一步是NLP中必做的一步，因为我们要根据输入的实体，找到对应的id，进而找到对应的embedding

生成训练数据

train_data_idxs=[[entity_idxs[triplet[0]],relation_idxs[triplet[1]],entity_idxs[triplet[2]]] for triplet in train_data]

生成批次的数据输入

er_vocab=defaultdict(list)
for triplet in train_data_idxs:
    er_vocab[(triplet[0],triplet[1])].append(triplet[2])
er_vocab_pairs=list(er_vocab.keys())

batch_inputs=er_vocab_pairs[:4]
batch_targets=torch.zeros([len(batch_inputs),len(entity_idxs)],dtype=torch.float32)
for i,pair in enumerate(batch_inputs):
    batch_targets[i,er_vocab[pair]]=1
batch_inputs=np.array(batch_inputs)

以前4个三元组为例

输入的是头实体和关系的id，输出的标签是尾实体的id。

需要注意的是，同一组头实体和关系，会有很多个尾实体存在的。

以第一个三元组为例：

即，(3920,791,9220)是一个三元组，也就是输入数据的一个样本。(3920,791,3799)也是一个三元组。所以标签有两个1。

了解了输入输出，接下来就可以定义模型。

模型

class KGE(nn.Module):
    def __init__(self,model_name,ent_vec_dim,num_entities,num_relations):
        '''
        num_entities是所有实体的数量
        num_relations是所有关系的数量
        ent_vec_dim是每一个实体向量的维度
        如果model_name是RESCAL，那么每一个关系用一个矩阵matrix表示，shape==(ent_vec_dim,ent_vec_dim)
        
        '''
        super(KGE,self).__init__()
        self.E=nn.Embedding(num_embeddings=num_entities,embedding_dim=ent_vec_dim,padding_idx=0)
        self.model_name=model_name
        self.ent_vec_dim=ent_vec_dim
        self.num_entities=num_entities
        if self.model_name=='RESCAL':
            self.R=nn.Embedding(num_embeddings=num_relations,embedding_dim=ent_vec_dim*ent_vec_dim,padding_idx=0)
            self.scoreFun=self.RESCAL
        else:
            self.R=nn.Embedding(num_embeddings=num_relations,embedding_dim=ent_vec_dim,padding_idx=0)
            self.scoreFun=self.DistMult
        
    def RESCAL(self,head_embed,rel_embed):
        '''
        RESCAL模型将每一个关系用一个matrix表示。
        输入：
            head_embed.size()==(batch_size,self.ent_vec_dim)
            rel_embed.size()==(batch_size,self.ent_vec_dim*2)
        输出：
            score.size()==(batch_size,self.num_entities)
        '''
        batch_size=head_embed.size(0)
        head_embed=head_embed.view(batch_size,1,self.ent_vec_dim)
        rel_embed=rel_embed.view(batch_size,self.ent_vec_dim,self.ent_vec_dim)
        score=torch.mm(torch.squeeze(torch.bmm(head_embed,rel_embed),dim=1),self.E.weight.transpose(1,0))
        return score
    
    def DistMult(self,head_embed,rel_embed):
        '''
        DistMult是RESCAL的简化版，将每一个关系用一个vector表示。
        输入：
            head_embed.size()==(batch_size,self.ent_vec_dim)
            rel_embed.size()==(batch_size,self.ent_vec_dim)
        输出：
            score.size()==(batch_size,self.num_entities)        
        '''
        score=torch.mm(head_embed*rel_embed,self.E.weight.transpose(1,0))
        return score
    
    def forward(self,head_idx,rel_idx):
        '''
        输入：
            head_idx.size()==rel_idx.size()==(batch_size,)
        输出：
            probabilities.size()==(batch_size,self.num_entities)     
            即：预测每一个实体可以作为尾实体的概率
        '''
        batch_size=head_idx.size(0)
        score=self.scoreFun(head_embed=self.E(head_idx),rel_embed=self.R(rel_idx))
        assert score.size()==(batch_size,self.num_entities)
        probabilities=torch.sigmoid(score)
        return probabilities

前向传播

head_idx=torch.LongTensor(batch_inputs[:,0])
rel_idx=torch.LongTensor(batch_inputs[:,1])

RESCAL=KGE(model_name='RESCAL',ent_vec_dim=200,num_entities=len(entity_idxs),num_relations=len(relation_idxs))
DistMult=KGE(model_name='DistMult',ent_vec_dim=200,num_entities=len(entity_idxs),num_relations=len(relation_idxs))

probabilities1=RESCAL(head_idx,rel_idx)
probabilities2=DistMult(head_idx,rel_idx)

计算BCE损失

loss=torch.nn.BCELoss()(probabilities1,batch_targets)
print(loss.item())
loss.backward()

这里简单说下BCEloss的计算方式

bce_loss=torch.nn.BCELoss()
x=torch.tensor([[3.4,4.5,3.1],[2.4,1.2,1.1]])
x=torch.sigmoid(x)
print(x)
y=torch.Tensor([[0,1,0],[0,0,1]])
print(bce_loss(x,y))


log=torch.log
a=log(torch.tensor(1-0.9677))+log(torch.tensor(0.9890))+log(torch.tensor(1-0.9569))
b=log(torch.tensor(1-0.9168))+log(torch.tensor(1-0.7685))+log(torch.tensor(0.7503))
print(-(a+b)/6)

所以计算公式为：
$$-\frac{1}{n\ast m}\sum _{i=1}^n\sum _{j=1}^m[y_j^{(i)}\ast \log p_j^{(i)}+(1-y_j^{(i)})\ast \log (1-p_j^{(i)})]$$

测试模型

假设模型已经训练好了，下面来测试模型

加载测试数据

test_data_idxs=[[entity_idxs[triplet[0]],relation_idxs[triplet[1]],entity_idxs[triplet[2]]] for triplet in test_data]
test_er_vocab=defaultdict(list)
for triplet in test_data_idxs:
    test_er_vocab[(triplet[0],triplet[1])].append(triplet[2])
    
test_batch_inputs=test_data_idxs[:3]
test_batch_inputs=np.array(test_batch_inputs)
test_batch_inputs

前向传播获取预测分数

head_idx=torch.tensor(test_batch_inputs[:,0])
rel_idx=torch.tensor(test_batch_inputs[:,1])
tail_idx=torch.tensor(test_batch_inputs[:,2])
probabilities=RESCAL(head_idx,rel_idx)

tail_idx是标签,probabilities是预测的分数

特别说明

测试评估阶段与训练阶段不同。

我们以测试数据中第一个样本为例，输入的头实体id是2431，关系id是89，尾实体id是5452，也就是(2431,89,5452)这么一个三元组，我们希望的是模型预测第5452个实体作为尾实体的概率越大越好。但是给定头实体2431和关系89，可不是只有一个三元组的

可以看到，在测试集中有三个三元组，即：
(2431,89,5452)
(2431,89,10961)
(2431,89,7741)
模型预测的probabilities的每一个数值代表的含义是哪个实体与当前给定的头实体和关系构成事实三元组的概率。因此模型如果预测第10961或者第7741个实体的分数大，那么并没有错。但是测试评估下，对于当前样本(2431,89,5452)，我们只关心模型预测第5452个实体的分数，而希望预测其他实体的分数越低越好。因此对于模型输出的probabilities，我们需要手动将第10961和7741这两个位置的数值置为0

for i in range(len(test_batch_inputs)):
    head,rel,tail=head_idx[i].item(),rel_idx[i].item(),tail_idx[i].item()
    #(head,rel,tail)是当前的三元组
    all_fact_tails=test_er_vocab[(head,rel)]#给定当前头实体和关系下，测试集中所有符合的尾实体
    print(head,rel,all_fact_tails)
    predict_score=probabilities[i][tail].item()#首先取出模型预测当前三元组尾实体的分数
    probabilities[i][all_fact_tails]=0.0#将测试集中所有与(head,rel,?)满足事实三元组的尾实体置空
    probabilities[i][tail]=predict_score#恢复模型预测当前三元组尾实体的分数

计算Mean Rank和hit@1,3,10

MeanRank: 模型输出的probabilities的长度是num_entities，我们将probabilities降序排列，即分数高的排在前面。mean rank中的rank指的就是：模型对于当前三元组尾实体在所有实体中的分数排名。
所以这个数值越低越好，因为越低，表明排名越靠前。

Hit@1,3,10: @就是英文的at。hit at 1指的就是将尾实体排在第一位的次数/测试集合大小，hit at 3和git at 10同理。显然这个数值越高越好。越高说明每一个三元组的尾实体都被模型排名的非常靠前。

sort_scores,sort_idxs=torch.sort(probabilities,dim=1,descending=True)
ranks=[]
hits=[[] for _ in range(10)]
for i in range(len(test_batch_inputs)):
    rank=np.where(sort_idxs[i].numpy()==tail_idx[i].item())[0][0]+1
    ranks.append(rank)
    
    for hits_level in range(1,11):
        if rank<=hits_level:
            hits[hits_level-1].append(1.0)
        else:
            hits[hits_level-1].append(0.0)

上图就是这三个测试样本的每一个的排名

hitat10=np.mean(hits[9])
hitat3=np.mean(hits[2])
hitat1=np.mean(hits[0])
mean_rank=np.mean(ranks)
mrr=np.mean(1./np.array(mean_rank))