class Train(nn.Module): def __init__(self,args,dataset): super(Train, self).__init__() self.args = args self.D = dataset self.entity_vec = nn.Embedding(self.D.entity_num,args.emb_dim) self.concept_vec = nn.Embedding(self.D.concept_num,args.emb_dim+1) self.relation_vec = nn.Embedding(self.D.relation_num,args.emb_dim) self.optimizer = torch.optim.SGD(self.parameters(),lr=args.lr) nn.init.normal_(self.entity_vec.weight.data, 0.0, 1.0 / args.emb_dim) nn.init.normal_(self.relation_vec.weight.data, 0.0, 1.0 / args.emb_dim) nn.init.normal_(self.concept_vec.weight.data[:, :-1], 0.0, 1.0 / args.emb_dim) nn.init.uniform_(self.concept_vec.weight.data[:, -1], 0.0, 1.0)

这是一个使用 PyTorch 实现的知识图谱嵌入模型，包含三个嵌入层：entity_vec、concept_vec 和 relation_vec。使用 nn.Embedding 可以将实体、概念和关系映射到低维空间中的向量表示。其中，entity_vec 和 relation_vec 的权重矩阵使用正态分布进行初始化，而 concept_vec 的权重矩阵的前 n-1 列也使用正态分布进行初始化，最后一列使用均匀分布进行初始化。该模型还定义了一个 SGD 优化器用于训练模型。此外，该模型还包含一个 self.D 参数，用于传入数据集。