import torch
import torch.nn as nn
import torch.nn.functional as F
from .Model import Model
class TransE(Model):
def __init__(self, ent_tot, rel_tot, dim=100, p_norm=1, norm_flag=True, margin=None, epsilon=None):
super(TransE, self).__init__(ent_tot, rel_tot) # 调用父类的初始化方法,ent_tot表示实体的总数,rel_tot表示关系的总数
# ent_tot代表实体的总数,rel_tot代表关系的总数。
self.dim = dim # 实体和关系的嵌入维度,默认为100
self.margin = margin # 损失函数的边界,默认为None
self.epsilon = epsilon # 嵌入的范围,默认为None
self.norm_flag = norm_flag # 是否对嵌入进行归一化,默认为True
self.p_norm = p_norm # 范数的大小,默认为1
self.ent_embeddings = nn.Embedding(self.ent_tot, self.dim) # 创建一个实体嵌入的Embedding层,大小为(ent_tot, dim)
self.rel_embeddings = nn.Embedding(self.rel_tot, self.dim) # 创建一个关系嵌入的Embedding层,大小为(rel_tot, dim)
if margin == None or epsilon == None: # 如果边界或范围为None
nn.init.xavier_uniform_(self.ent_embeddings.weight.data) # 初始化实体嵌入的权重矩阵
nn.init.xavier_uniform_(self.rel_embeddings.weight.data) # 初始化关系嵌入的权重矩阵
else: # 如果边界和范围都不为None
self.embedding_range = nn.Parameter(
torch.Tensor([(self.margin + self.epsilon) / self.dim]), requires_grad=False
) # 创建一个参数embedding_range,用于限制嵌入的范围
nn.init.uniform_(
tensor=self.ent_embeddings.weight.data,
a=-self.embedding_range.item(),
b=self.embedding_range.item()
) # 初始化实体嵌入的权重矩阵,范围为[-embedding_range.item(), embedding_range.item()]
nn.init.uniform_(
tensor=self.rel_embeddings.weight.data,
a=-self.embedding_range.item(),
b=self.embedding_range.item()
) # 初始化关系嵌入的权重矩阵,范围为[-embedding_range.item(), embedding_range.item()]
if margin != None: # 如果边界不为None
self.margin = nn.Parameter(torch.Tensor([margin])) # 创建一个参数margin,用于限制损失函数的边界
self.margin.requires_grad = False # 设置margin参数不需要梯度计算
self.margin_flag = True # 设置margin_flag为True,表示使用了margin边界
else: # 如果边界为None
self.margin_flag = False # 设置margin_flag为False,表示没有使用margin边界
def _calc(self, h, t, r, mode):
if self.norm_flag: # 如果需要对嵌入进行归一化
h = F.normalize(h, 2, -1) # 对实体嵌入h进行L2范数归一化
r = F.normalize(r, 2, -1) # 对关系嵌入r进行L2范数归一化
t = F.normalize(t, 2, -1) # 对实体嵌入t进行L2范数归一化
if mode != 'normal': # 如果不是normal模式
h = h.view(-1, r.shape[0], h.shape[-1]) # 将实体嵌入h调整为(batch_size, relation_count, embedding_dim)的形状
t = t.view(-1, r.shape[0], t.shape[-1]) # 将实体嵌入t调整为(batch_size, relation_count, embedding_dim)的形状
r = r.view(-1, r.shape[0], r.shape[-1]) # 将关系嵌入r调整为(batch_size, relation_count, embedding_dim)的形状
if mode == 'head_batch': # 如果是head_batch模式
score = h + (r - t) # 计算得分,对应于公式h + (r - t)
else: # 如果是其他模式(关系批量或normal)
score = (h + r) - t # 计算得分,对应于公式(h + r) - t
score = torch.norm(score, self.p_norm, -1).flatten() # 计算得分的Lp范数,并展平为一维向量
return score # 返回计算得分结果
def forward(self, data):
batch_h = data['batch_h'] # 从输入数据中获取batch中的头实体索引,batch_h是一个张量
batch_t = data['batch_t'] # 从输入数据中获取batch中的尾实体索引,batch_t是一个张量
batch_r = data['batch_r'] # 从输入数据中获取batch中的关系索引,batch_r是一个张量
mode = data['mode'] # 从输入数据中获取计算模式,mode是一个字符串,可能的取值为'head_batch'、'tail_batch'或'normal'
h = self.ent_embeddings(batch_h) # 获取batch中头实体的嵌入向量h
t = self.ent_embeddings(batch_t) # 获取batch中尾实体的嵌入向量t
r = self.rel_embeddings(batch_r) # 获取batch中关系的嵌入向量r
score = self._calc(h, t, r, mode) # 调用_calc方法计算得分,根据不同的模式计算不同的得分
if self.margin_flag: # 如果使用了边界
return self.margin - score # 返回边界减去得分作为最终的损失值
else: # 如果没有使用边界
return score # 返回得分作为最终的损失值
def regularization(self, data):
batch_h = data['batch_h'] # 从输入数据中获取batch中的头实体索引,batch_h是一个张量
batch_t = data['batch_t'] # 从输入数据中获取batch中的尾实体索引,batch_t是一个张量
batch_r = data['batch_r'] # 从输入数据中获取batch中的关系索引,batch_r是一个张量
h = self.ent_embeddings(batch_h) # 获取batch中头实体的嵌入向量h
t = self.ent_embeddings(batch_t) # 获取batch中尾实体的嵌入向量t
r = self.rel_embeddings(batch_r) # 获取batch中关系的嵌入向量r
regul = (torch.mean(h ** 2) + # 计算头实体嵌入的L2范数的平均值
torch.mean(t ** 2) + # 计算尾实体嵌入的L2范数的平均值
torch.mean(r ** 2)) / 3 # 计算关系嵌入的L2范数的平均值,然后取三者的平均值作为正则化项
return regul # 返回正则化项
def predict(self, data): #测试,不用管
score = self.forward(data) # 调用forward方法预测得分
if self.margin_flag: # 如果使用了边界
score = self.margin - score # 计算边界减去得分作为最终得分
return score.cpu().data.numpy() # 将最终得分转换为numpy数组并返回
else: # 如果没有使用边界
return score.cpu().data.numpy() # 将得分转换为numpy数组并返回