unike 3.0.1__py3-none-any.whl

unike/module/model/RGCN.py

@@ -0,0 +1,304 @@
+# coding:utf-8
+#
+# unike/module/model/RGCN.py
+# 
+# created by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on Jan 16, 2024
+# updated by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on May 22, 2024
+# 
+# 该头文件定义了 R-GCN.
+
+"""
+R-GCN - 第一个图神经网络模型。
+"""
+
+import dgl
+import torch
+import typing
+import torch.nn as nn
+import torch.nn.functional as F
+from .Model import Model
+from dgl.nn.pytorch import RelGraphConv
+from typing_extensions import override
+
+class RGCN(Model):
+    
+    """
+	``R-GCN`` :cite:`R-GCN` 提出于 2017 年，是第一个图神经网络模型。
+	
+	正三元组的评分函数的值越大越好，如果想获得更详细的信息请访问 :ref:`R-GCN <rgcn>`。
+
+    例子::
+
+        from unike.data import GraphDataLoader
+        from unike.module.model import RGCN
+        from unike.module.loss import RGCNLoss
+        from unike.module.strategy import RGCNSampling
+        from unike.config import Trainer, GraphTester
+        
+        dataloader = GraphDataLoader(
+        	in_path = "../../benchmarks/FB15K237/",
+        	batch_size = 60000,
+        	neg_ent = 10,
+        	test = True,
+        	test_batch_size = 100,
+        	num_workers = 16
+        )
+        
+        # define the model
+        rgcn = RGCN(
+        	ent_tol = dataloader.get_ent_tol(),
+        	rel_tol = dataloader.get_rel_tol(),
+        	dim = 500,
+        	num_layers = 2
+        )
+        
+        # define the loss function
+        model = RGCNSampling(
+        	model = rgcn,
+        	loss = RGCNLoss(model = rgcn, regularization = 1e-5)
+        )
+        
+        # test the model
+        tester = GraphTester(model = rgcn, data_loader = dataloader, use_gpu = True, device = 'cuda:0')
+        
+        # train the model
+        trainer = Trainer(model = model, data_loader = dataloader.train_dataloader(),
+        	epochs = 10000, lr = 0.0001, use_gpu = True, device = 'cuda:0',
+        	tester = tester, test = True, valid_interval = 500, log_interval = 500,
+        	save_interval = 500, save_path = '../../checkpoint/rgcn.pth'
+        )
+        trainer.run()
+	"""
+
+    def __init__(
+        self,
+        ent_tol: int,
+        rel_tol: int,
+        dim: int,
+        num_layers: int):
+
+        """创建 RGCN 对象。
+
+		:param ent_tol: 实体的个数
+		:type ent_tol: int
+		:param rel_tol: 关系的个数
+		:type rel_tol: int
+		:param dim: 实体和关系嵌入向量的维度
+		:type dim: int
+		:param num_layers: 图神经网络的层数
+		:type num_layers: int
+		"""
+
+        super(RGCN, self).__init__(ent_tol, rel_tol)
+
+        #: 实体和关系嵌入向量的维度
+        self.dim: int = dim
+        #: 图神经网络的层数
+        self.num_layers: int = num_layers
+
+        #: 根据实体个数，创建的实体嵌入
+        self.ent_emb: torch.nn.Embedding = None
+        #: 根据关系个数，创建的关系嵌入
+        self.rel_emb: torch.nn.parameter.Parameter = None
+        #: R-GCN 的图神经网络层
+        self.RGCN: torch.nn.ModuleList = None
+        #: 图神经网络层的输出
+        self.Loss_emb: torch.nn.Embedding = None
+
+        self.build_model()
+
+    def build_model(self):
+
+        """构建模型"""
+
+        self.ent_emb = nn.Embedding(self.ent_tol, self.dim)
+
+        self.rel_emb = nn.Parameter(torch.Tensor(self.rel_tol, self.dim))
+
+        nn.init.xavier_uniform_(self.rel_emb, gain=nn.init.calculate_gain('relu'))
+
+        self.RGCN = nn.ModuleList()
+        for idx in range(self.num_layers):
+            RGCN_idx = self.build_hidden_layer(idx)
+            self.RGCN.append(RGCN_idx)
+
+    def build_hidden_layer(
+        self,
+        idx: int) -> dgl.nn.pytorch.conv.RelGraphConv:
+
+        """返回第 idx 的图神经网络层。
+        
+        :param idx: 数据。
+        :type idx: int
+        :returns: 图神经网络层
+        :rtype: dgl.nn.pytorch.conv.RelGraphConv
+        """
+
+        act = F.relu if idx < self.num_layers - 1 else None
+        return RelGraphConv(self.dim, self.dim, self.rel_tol, "bdd",
+                    num_bases=100, activation=act, self_loop=True, dropout=0.2)
+
+    @override  
+    def forward(
+        self,
+        graph: dgl.DGLGraph,
+        ent: torch.Tensor,
+        rel: torch.Tensor,
+        norm: torch.Tensor,
+        triples: torch.Tensor,
+        mode: str = 'single') -> torch.Tensor:
+
+        """
+		定义每次调用时执行的计算。
+		:py:class:`torch.nn.Module` 子类必须重写 :py:meth:`torch.nn.Module.forward`。
+        
+        :param graph: 子图
+        :type graph: dgl.DGLGraph
+        :param ent: 子图的实体
+        :type ent: torch.Tensor
+        :param rel: 子图的关系
+        :type rel: torch.Tensor
+        :param norm: 关系的归一化系数
+        :type norm: torch.Tensor
+        :param triples: 三元组
+        :type triples: torch.Tensor
+        :param mode: 模式
+        :type mode: str
+        :returns: 三元组的得分
+        :rtype: torch.Tensor
+		"""
+        
+        embedding = self.ent_emb(ent.squeeze())
+        for layer in self.RGCN:
+            embedding = layer(graph, embedding, rel, norm)
+        self.Loss_emb = embedding
+        head_emb, rela_emb, tail_emb = self.tri2emb(embedding, triples, mode)
+        score = self.distmult_score_func(head_emb, rela_emb, tail_emb, mode)
+
+        return score
+
+    def tri2emb(
+        self,
+        embedding: torch.Tensor,
+        triples: torch.Tensor,
+        mode: str = "single") -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+
+        """
+		获得三元组对应头实体、关系和尾实体的嵌入向量。
+        
+        :param embedding: 经过图神经网络更新的实体嵌入向量
+        :type embedding: torch.Tensor
+        :param triples: 训练的三元组
+        :type triples: torch.Tensor
+        :param mode: 模式
+        :type mode: str
+        :returns: 头实体、关系和尾实体的嵌入向量
+        :rtype: torch.Tensor
+		"""
+
+        rela_emb = self.rel_emb[triples[:, 1]].unsqueeze(1)  # [bs, 1, dim]
+        head_emb = embedding[triples[:, 0]].unsqueeze(1)  # [bs, 1, dim] 
+        tail_emb = embedding[triples[:, 2]].unsqueeze(1)  # [bs, 1, dim]
+
+        if mode == "head-batch" or mode == "head_predict":
+            head_emb = embedding.unsqueeze(0)  # [1, num_ent, dim]
+
+        elif mode == "tail-batch" or mode == "tail_predict":
+            tail_emb = embedding.unsqueeze(0)  # [1, num_ent, dim]
+
+        return head_emb, rela_emb, tail_emb
+
+    def distmult_score_func(
+        self,
+        head_emb: torch.Tensor,
+        relation_emb: torch.Tensor,
+        tail_emb: torch.Tensor,
+        mode: str) -> torch.Tensor:
+
+        """
+		计算 DistMult 的评分函数。
+        
+        :param head_emb: 头实体嵌入向量
+        :type head_emb: torch.Tensor
+        :param relation_emb: 关系嵌入向量
+        :type relation_emb: torch.Tensor
+        :param tail_emb: 尾实体嵌入向量
+        :type tail_emb: torch.Tensor
+        :returns: 三元组的得分
+        :rtype: torch.Tensor
+		"""
+
+        if mode == 'head-batch':
+            score = head_emb * (relation_emb * tail_emb)
+        else:
+            score = (head_emb * relation_emb) * tail_emb
+
+        score = score.sum(dim = -1)
+        return score
+
+    @override
+    def predict(
+        self,
+        data: dict[str, torch.Tensor],
+        mode: str) -> torch.Tensor:
+
+        """R-GCN 的推理方法。
+        
+        :param data: 数据。
+        :type data: dict[str, torch.Tensor]
+        :param mode: 模式
+        :type mode: str
+        :returns: 三元组的得分
+        :rtype: torch.Tensor
+		"""
+
+        triples    = data['positive_sample']
+        graph      = data['graph']
+        ent        = data['entity']
+        rel        = data['rela']
+        norm       = data['norm']
+
+        embedding = self.ent_emb(ent.squeeze())
+        for layer in self.RGCN:
+            embedding = layer(graph, embedding, rel, norm)
+        self.Loss_emb = embedding
+        head_emb, rela_emb, tail_emb = self.tri2emb(embedding, triples, mode)
+        score = self.distmult_score_func(head_emb, rela_emb, tail_emb, mode)
+
+        return score
+
+def get_rgcn_hpo_config() -> dict[str, dict[str, typing.Any]]:
+
+	"""返回 :py:class:`RGCN` 的默认超参数优化配置。
+	
+	默认配置为::
+	
+	    parameters_dict = {
+	    	'model': {
+	    		'value': 'RGCN'
+	    	},
+	    	'dim': {
+	    		'values': [200, 300, 400]
+	    	},
+	    	'num_layers': {
+	    		'value': 2
+	    	}
+	    }
+
+	:returns: :py:class:`RGCN` 的默认超参数优化配置
+	:rtype: dict[str, dict[str, typing.Any]]
+	"""
+
+	parameters_dict = {
+		'model': {
+			'value': 'RGCN'
+		},
+		'dim': {
+			'values': [200, 300, 400]
+		},
+		'num_layers': {
+			'value': 2
+		}
+	}
+		
+	return parameters_dict

unike/module/model/RotatE.py

@@ -0,0 +1,303 @@
+# coding:utf-8
+#
+# unike/module/model/RotatE.py
+# 
+# git pull from OpenKE-PyTorch by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on May 7, 2023
+# updated by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on Jan 11, 2023
+# 
+# 该头文件定义了 RotatE.
+
+"""
+RotatE - 将实体表示成复数向量，关系建模为复数向量空间的旋转。
+"""
+
+import torch
+import typing
+import torch.nn as nn
+from .Model import Model
+from typing_extensions import override
+
+class RotatE(Model):
+
+	"""
+	``RotatE`` :cite:`RotatE` 提出于 2019 年，将实体表示成复数向量，关系建模为复数向量空间的旋转。
+
+	评分函数为:
+
+	.. math::
+	
+		\gamma - \parallel \mathbf{h} \circ \mathbf{r} - \mathbf{t} \parallel_{L_2}
+	
+	:math:`\circ` 表示哈达玛积（Hadamard product），正三元组的评分函数的值越大越好，如果想获得更详细的信息请访问 :ref:`RotatE <rotate>`。
+
+	例子::
+
+		from unike.data import KGEDataLoader, UniSampler, TradTestSampler
+		from unike.module.model import RotatE
+		from unike.module.loss import SigmoidLoss
+		from unike.module.strategy import NegativeSampling
+		from unike.config import Trainer, Tester
+		
+		# dataloader for training
+		dataloader = KGEDataLoader(
+			in_path = '../../benchmarks/WN18RR/', 
+			batch_size = 2000,
+			neg_ent = 64,
+			test = True,
+			test_batch_size = 10,
+			num_workers = 16,
+			train_sampler = UniSampler,
+			test_sampler = TradTestSampler
+		)
+		
+		# define the model
+		rotate = RotatE(
+			ent_tol = dataloader.get_ent_tol(),
+			rel_tol = dataloader.get_rel_tol(),
+			dim = 1024,
+			margin = 6.0,
+			epsilon = 2.0,
+		)
+		
+		# define the loss function
+		model = NegativeSampling(
+			model = rotate, 
+			loss = SigmoidLoss(adv_temperature = 2), 
+			regul_rate = 0.0,
+		)
+		
+		# test the model
+		tester = Tester(model = rotate, data_loader = dataloader, use_gpu = True, device = 'cuda:1')
+		
+		# train the model
+		trainer = Trainer(model = model, data_loader = dataloader.train_dataloader(), epochs = 6000,
+			lr = 2e-5, opt_method = 'adam', use_gpu = True, device = 'cuda:1',
+			tester = tester, test = True, valid_interval = 100,
+			log_interval = 100, save_interval = 100,
+			save_path = '../../checkpoint/rotate.pth', use_wandb = False)
+		trainer.run()
+	"""
+
+	def __init__(
+		self,
+		ent_tol: int,
+		rel_tol: int,
+		dim: int = 100,
+		margin: float = 6.0,
+		epsilon: float = 2.0):
+
+		"""创建 RotatE 对象。
+
+		:param ent_tol: 实体的个数
+		:type ent_tol: int
+		:param rel_tol: 关系的个数
+		:type rel_tol: int
+		:param dim: 实体和关系嵌入向量的维度
+		:type dim: int
+		:param margin: 原论文中损失函数的 gamma。
+		:type margin: float
+		:param epsilon: RotatE 原论文对应的源代码固定为 2.0。
+		:type epsilon: float
+		"""
+
+		super(RotatE, self).__init__(ent_tol, rel_tol)
+
+		#: RotatE 原论文对应的源代码固定为 2.0。
+		self.epsilon: int = epsilon
+
+		#: RotatE 原论文的实现中将实体嵌入向量的维度指定为 ``dim`` 的 2 倍。
+		#: 因为实体嵌入向量需要划分为实部和虚部。
+		self.dim_e: int = dim * 2
+		#: 关系嵌入向量的维度，为 ``dim``。
+		self.dim_r: int = dim
+
+		#: 根据实体个数，创建的实体嵌入。
+		self.ent_embeddings: torch.nn.Embedding = nn.Embedding(self.ent_tol, self.dim_e)
+		#: 根据关系个数，创建的关系嵌入。
+		self.rel_embeddings: torch.nn.Embedding = nn.Embedding(self.rel_tol, self.dim_r)
+
+		self.ent_embedding_range = nn.Parameter(
+			torch.Tensor([(margin + self.epsilon) / self.dim_e]), 
+			requires_grad=False
+		)
+
+		nn.init.uniform_(
+			tensor = self.ent_embeddings.weight.data, 
+			a=-self.ent_embedding_range.item(), 
+			b=self.ent_embedding_range.item()
+		)
+
+		self.rel_embedding_range = nn.Parameter(
+			torch.Tensor([(margin + self.epsilon) / self.dim_r]), 
+			requires_grad=False
+		)
+
+		nn.init.uniform_(
+			tensor = self.rel_embeddings.weight.data, 
+			a=-self.rel_embedding_range.item(), 
+			b=self.rel_embedding_range.item()
+		)
+
+		#: 原论文中损失函数的 gamma。
+		self.margin: torch.nn.parameter.Parameter = nn.Parameter(torch.Tensor([margin]))
+		self.margin.requires_grad = False
+
+	@override
+	def forward(
+		self,
+		triples: torch.Tensor,
+		negs: torch.Tensor = None,
+		mode: str = 'single') -> torch.Tensor:
+
+		"""
+		定义每次调用时执行的计算。
+		:py:class:`torch.nn.Module` 子类必须重写 :py:meth:`torch.nn.Module.forward`。
+		
+		:param triples: 正确的三元组
+		:type triples: torch.Tensor
+		:param negs: 负三元组类别
+		:type negs: torch.Tensor
+		:param mode: 模式
+		:type triples: str
+		:returns: 三元组的得分
+		:rtype: torch.Tensor
+		"""
+
+		head_emb, relation_emb, tail_emb = self.tri2emb(triples, negs, mode)
+		score = self.margin - self._calc(head_emb, relation_emb, tail_emb)
+		return score
+
+	def _calc(
+		self,
+		h: torch.Tensor,
+		r: torch.Tensor,
+		t: torch.Tensor) -> torch.Tensor:
+
+		"""计算 RotatE 的评分函数。
+
+		利用 :py:func:`torch.chunk` 拆分实体嵌入向量获得复数的实部和虚部。
+		原论文使用 L1-norm 作为距离函数，而这里使用的 L2-norm 作为距离函数。
+		
+		:param h: 头实体的向量。
+		:type h: torch.Tensor
+		:param r: 关系的向量。
+		:type r: torch.Tensor
+		:param t: 尾实体的向量。
+		:type t: torch.Tensor
+		:returns: 三元组的得分
+		:rtype: torch.Tensor
+		"""
+
+		pi = self.pi_const
+
+		re_head, im_head = torch.chunk(h, 2, dim=-1)
+		re_tail, im_tail = torch.chunk(t, 2, dim=-1)
+
+		# Make phases of relations uniformly distributed in [-pi, pi]
+		phase_relation = r / (self.rel_embedding_range.item() / pi)
+
+		re_relation = torch.cos(phase_relation)
+		im_relation = torch.sin(phase_relation)
+	
+		re_score = re_head * re_relation - im_head * im_relation
+		im_score = re_head * im_relation + im_head * re_relation
+		re_score = re_score - re_tail
+		im_score = im_score - im_tail
+
+		score = torch.stack([re_score, im_score], dim = 0)
+		score = score.norm(dim = 0).sum(dim = -1)
+		return score
+
+	@override
+	def predict(
+		self,
+		data: dict[str, typing.Union[torch.Tensor,str]],
+		mode) -> torch.Tensor:
+		
+		"""RotatE 的推理方法。
+		
+		:param data: 数据。
+		:type data: dict[str, typing.Union[torch.Tensor,str]]
+		:returns: 三元组的得分
+		:rtype: torch.Tensor
+		"""
+
+		triples = data["positive_sample"]
+		head_emb, relation_emb, tail_emb = self.tri2emb(triples, mode=mode)
+		score = self.margin - self._calc(head_emb, relation_emb, tail_emb)
+		return score
+
+	def regularization(
+		self,
+		data: dict[str, typing.Union[torch.Tensor, str]]) -> torch.Tensor:
+
+		"""L2 正则化函数（又称权重衰减），在损失函数中用到。
+		
+		:param data: 数据。
+		:type data: dict[str, typing.Union[torch.Tensor, str]]
+		:returns: 模型参数的正则损失
+		:rtype: torch.Tensor
+		"""
+
+		pos_sample = data["positive_sample"]
+		neg_sample = data["negative_sample"]
+		mode = data["mode"]
+		pos_head_emb, pos_relation_emb, pos_tail_emb = self.tri2emb(pos_sample)
+		if mode == "bern":
+			neg_head_emb, neg_relation_emb, neg_tail_emb = self.tri2emb(neg_sample)
+		else:
+			neg_head_emb, neg_relation_emb, neg_tail_emb = self.tri2emb(pos_sample, neg_sample, mode)
+
+		pos_regul = (torch.mean(pos_head_emb ** 2) + 
+					 torch.mean(pos_relation_emb ** 2) + 
+					 torch.mean(pos_tail_emb ** 2)) / 3
+
+		neg_regul = (torch.mean(neg_head_emb ** 2) + 
+					 torch.mean(neg_relation_emb ** 2) + 
+					 torch.mean(neg_tail_emb ** 2)) / 3
+
+		regul = (pos_regul + neg_regul) / 2
+
+		return regul
+
+def get_rotate_hpo_config() -> dict[str, dict[str, typing.Any]]:
+
+	"""返回 :py:class:`RotatE` 的默认超参数优化配置。
+	
+	默认配置为::
+	
+		parameters_dict = {
+			'model': {
+				'value': 'RotatE'
+			},
+			'dim': {
+				'values': [256, 512, 1024]
+			},
+			'margin': {
+				'values': [1.0, 3.0, 6.0]
+			},
+			'epsilon': {
+				'value': 2.0
+			}
+		}
+
+	:returns: :py:class:`RotatE` 的默认超参数优化配置
+	:rtype: dict[str, dict[str, typing.Any]]
+	"""
+
+	parameters_dict = {
+		'model': {
+			'value': 'RotatE'
+		},
+		'dim': {
+			'values': [256, 512, 1024]
+		},
+		'margin': {
+			'values': [1.0, 3.0, 6.0]
+		},
+		'epsilon': {
+			'value': 2.0
+		}
+	}
+		
+	return parameters_dict