openKE-Transe

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数组并返回