简单实现几篇知识图谱嵌入(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 c o u n t r y − V c i t y V_{country}-V_{city}V c o u n t r y ​−V c i t y ​ 得到的向量与V c a p t i a l − o f V_{captial-of}V c a p t i a l −o f ​ 这个向量就很相近。TransE也认为,在KG的embedding空间中,也一定存在这种关系,即两个实体的嵌入向量的差值(或者其它操作)就代表这两个实体之间的关系

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

数学定义中:
平移前点的坐标+平移向量的坐标=平移后点的坐标

所以我们称之为 平移

TransE的得分函数就是向量之间的欧氏距离的相反数:
S c o r e ( h , r , t ) = − ∣ ∣ v h + v r − v t ∣ ∣ 2 2 Score(h,r,t)=-||v_h+v_r-v_t||_2^2 S c o r e (h ,r ,t )=−∣∣v h ​+v r ​−v t ​∣∣2 2 ​
损失函数定义为:
m a x ( 0 , γ − S c o r e ( h , r , t ) + S c o r e ( h ′ , r , t ′ ) ) max(0,\gamma-Score(h,r,t)+Score(h’,r,t’))m a x (0 ,γ−S c o r e (h ,r ,t )+S c o r e (h ′,r ,t ′))
也就是在embedding space中,正例三元组的得分要比负例三元组的得分高出γ \gamma γ,又由于得分函数表示为距离的相反数,所以得分高代表距离近。即:正例三元组的距离要比负例三元组的距离小至少γ \gamma γ长度的距离。

RESCAL

该模型的 得分函数定义为:
S c o r e ( h , r , t ) = v h T M r v t Score(h,r,t)=v_h^TM_rv_t S c o r e (h ,r ,t )=v h T ​M r ​v t ​

其中v h ∈ R d , v t ∈ R d , M r ∈ R d × d v_h\in R^d,v_t\in R^d,M_r\in R^{d\times d}v h ​∈R d ,v t ​∈R d ,M r ​∈R d ×d
v h , v t v_h,v_t v h ​,v t ​都是从实体embedding矩阵(记为E E E)中的 取出(根据实体id获取)的vector,E E E的形状是(num_entities,d)。
M r M_r M r ​是整个关系tensor(三维的)中的根据关系id获取的matrix(二维的)。整个关系tensor记为R R R,形状是(num_relations,d,d),所以M r M_r M r ​的shape是(d , d d,d 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 ∈ R d v_h,v_r,v_t\in R^d v h ​,v r ​,v t ​∈R d

代码实现

获取数据

我们使用FB15K知识库

下载解压:

简单实现几篇知识图谱嵌入(Knowledge Graph Embedding,KGE)模型
FB15K知识库就是TransE这篇论文的作者从Freebase知识库中选取的一部分三元组构成的一个小规模的知识库

三元组数量实体数量关系数量592213149511345

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

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

即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]

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

生成批次的数据输入

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个三元组为例

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

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

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

以第一个三元组为例:

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

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

简单实现几篇知识图谱嵌入(Knowledge Graph Embedding,KGE)模型
了解了输入输出,接下来就可以定义模型。

模型

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])

简单实现几篇知识图谱嵌入(Knowledge Graph Embedding,KGE)模型
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()

简单实现几篇知识图谱嵌入(Knowledge Graph Embedding,KGE)模型
这里简单说下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)

简单实现几篇知识图谱嵌入(Knowledge Graph Embedding,KGE)模型
所以计算公式为:
− 1 n ∗ m ∑ i = 1 n ∑ j = 1 m [ y j ( i ) ∗ log ⁡ p j ( i ) + ( 1 − y j ( i ) ) ∗ log ⁡ ( 1 − p j ( i ) ) ] -\frac{1}{nm}\sum_{i=1}^{n}\sum_{j=1}^{m}[y_j^{(i)}\log p_j^{(i)}+(1-y_j^{(i)})*\log(1-p_j^{(i)})]−n ∗m 1 ​i =1 ∑n ​j =1 ∑m ​[y j (i )​∗lo g p j (i )​+(1 −y j (i )​)∗lo g (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

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

前向传播获取预测分数

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)

简单实现几篇知识图谱嵌入(Knowledge Graph Embedding,KGE)模型
tail_idx是标签,probabilities是预测的分数

特别说明

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

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

简单实现几篇知识图谱嵌入(Knowledge Graph Embedding,KGE)模型
可以看到,在测试集中有三个三元组,即:
(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()

    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
    probabilities[i][tail]=predict_score

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

计算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 rankhits_level:
            hits[hits_level-1].append(1.0)
        else:
            hits[hits_level-1].append(0.0)

简单实现几篇知识图谱嵌入(Knowledge Graph Embedding,KGE)模型
上图就是这三个测试样本的每一个的排名
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))

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

Original: https://blog.csdn.net/m0_45478865/article/details/121195480
Author: xhsun1997
Title: 简单实现几篇知识图谱嵌入(Knowledge Graph Embedding,KGE)模型

原创文章受到原创版权保护。转载请注明出处:https://www.johngo689.com/554717/

转载文章受原作者版权保护。转载请注明原作者出处!

(0)

大家都在看

亲爱的 Coder【最近整理,可免费获取】👉 最新必读书单  | 👏 面试题下载  | 🌎 免费的AI知识星球