Source code for models.DistMult

import torch
import torch.autograd as autograd
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
from Model import *
[docs]class DistMult(Model): r"""THis class implements DistMult model based on :class:`models.Model` Args: config: configuration information Shape: - Output: loss value which is a scalar """ def __init__(self,config): super(DistMult,self).__init__(config) self.ent_embeddings=nn.Embedding(self.config.entTotal,self.config.hidden_size) self.rel_embeddings=nn.Embedding(self.config.relTotal,self.config.hidden_size) self.softplus=nn.Softplus().cuda() self.init_weights()
[docs] def init_weights(self): r"""Init entity and relation embeddings. In this model, we use ``xaview_uniform`` """ nn.init.xavier_uniform( nn.init.xavier_uniform(
def _calc(self,h,t,r): return torch.sum(h*t*r,1,False)
[docs] def loss_func(self,loss,regul): r"""This method calculatest the regularized loss Args: loss (float): the loss without regularization regul (float): the regularization for loss Returns: float: the regularized loss """ return loss+self.config.lmbda*regul
def forward(self): batch_h,batch_t,batch_r=self.get_all_instance() batch_y=self.get_all_labels() e_h=self.ent_embeddings(batch_h) e_t=self.ent_embeddings(batch_t) e_r=self.rel_embeddings(batch_r) y=batch_y res=self._calc(e_h,e_t,e_r) tmp=self.softplus(- y * res) loss = torch.mean(tmp) regul = torch.mean(e_h ** 2) + torch.mean(e_t ** 2) + torch.mean(e_r ** 2) #Calculating loss to get what the framework will optimize loss = self.loss_func(loss,regul) return loss
[docs] def predict(self, predict_h, predict_t, predict_r): r"""This method predictss the score for testting sample Args: predict_h: head entity of triple predict_t: tail entity of triple predict_r: relaiton of triple Returns: float: the score for testing triple """ p_e_h=self.ent_embeddings(Variable(torch.from_numpy(predict_h)).cuda()) p_e_t=self.ent_embeddings(Variable(torch.from_numpy(predict_t)).cuda()) p_e_r=self.rel_embeddings(Variable(torch.from_numpy(predict_r)).cuda()) p_score=-self._calc(p_e_h,p_e_t,p_e_r) return p_score.cpu()