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

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

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

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

该模型将关系视为从头到尾实体的翻译。 TransE受到word2vec如果你已经训练了词向量，那么对于三个词：国家country、城市city，首都captial-of，会有以下关系: V c o u n t r y ? V c i t y V_{country}-V_{city} Vcountry?Vcity得到的向量和 V c a p t i a l ? o f V_{captial-of} Vcaptial?of这个向量非常相似。TransE也认为，在KG的embedding这种关系也必须存在于空间中，即两个实体的嵌入向量差（或其他操作）代表两个实体之间的关系。

由于关系和实体都被表示为向量，另一种数学说法是向量(vector)空间中，TransE将关系视为头实体到尾实体的平移操作，即： v h + v r ≈ v t v_h+v_r\approx v_t vh​+vr​≈vt​

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

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 Score(h,r,t)=−∣∣vh​+vr​−vt​∣∣22​ 损失函数定义为： 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')) max(0,γ−Score(h,r,t)+Score(h′,r,t′)) 也就是在embedding space中，正例三元组的得分要比负例三元组的得分高出 γ \gamma γ，又由于得分函数表示为距离的相反数，所以得分高代表距离近。即：正例三元组的距离要比负例三元组的距离小至少 γ \gamma γ长度的距离。

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

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

我们使用FB15K知识库

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

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

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

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

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个关系

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)