小黑笔记:transe模型

1.数据集准备

import openke
from openke.data import TrainDataLoader,TestDataLoader

train_dataloader = TrainDataLoader(
    in_path = "./benchmarks/FB15K237_tiny/",
    nbatches = 100,
    threads = 8,
    # 负采样
    sampling_mode = 'normal',
    # bern构建负样本方式
    bern_flag = 1,
    # 负样本同replace entities
    neg_ent = 25,
    neg_rel = 0
)
# dataloader for test
test_dataloader = TestDataLoader("./benchmarks/FB15K237_tiny/", "link")

2.模型构建

import torch
import torch.nn as nn
import torch.nn.functional as F
class TransE(nn.Module):
    def __init__(self,ent_tot,rel_tot,dim = 100,p_norm = 1,norm_flag = True,margin = None,epsilon = None):
        super(TransE,self).__init__()
        self.ent_tot = ent_tot
        self.rel_tot = rel_tot
        # embedding维度
        self.dim = dim
        self.margin = margin
        # 初始化为None
        self.epsilon = epsilon
        self.norm_flag = norm_flag
        # scoring function计算距离使用L1 norm
        self.p_norm = p_norm
        # embedding矩阵 [e,d]
        self.ent_embeddings = nn.Embedding(self.ent_tot,self.dim)
        # embedding矩阵,[r,d]
        self.rel_embeddings = nn.Embedding(self.rel_tot,self.dim)
        # 参数初始化
        if margin == None or epsilon == None:
            nn.init.xavier_uniform_(self.ent_embeddings.weight.data)
            nn.init.xavier_uniform_(self.rel_embeddings.weight.data)
        else:
            self.embedding_range = nn.Parameter(
                torch.Tensor([(self.margin + self.epsilon) / self.dim]),requires_grad = False
            )
            nn.init.uniform_(
                tensor = self.ent_embeddings.weight.data,
                a = -self.embedding_range.item(),
                b = self.embedding_range.item()
            )
            nn.init.uniform_(
                tensor = self.rel_embeddings.weight.data,
                a = -self.embedding_range.item(),
                b = self.embedding_range.item()
            )
        # 容忍系数margin
        if margin != None:
            self.margin = nn.Parameter(torch.Tensor([margin]))
            self.margin.requires_grad = False
            self.margin_flag = True
        else:
            # 执行
            self.margin_flag = False
    # 计算得分函数
    def _calc(self,h,t,r,mode):
        if self.norm_flag:
            h = F.normalize(h,2,-1)
            r = F.normalize(r,2,-1)
            t = F.normalize(t,2,-1)
        if mode != 'normal':
            h = h.view(-1,r.shape[0],h.shape[-1])
            t = t.view(-1,r.shape[0],t.shape[-1])
            r = r.view(-1,r.shape[0],r.shape[-1])
        if mode == 'head_batch':
            score = h + (r - t)
        else:
            score = (h + r) - t
        score = torch.norm(score,self.p_norm,-1).flatten()
        return score
    def forward(self,data):
        batch_h = data['batch_h']
        batch_t = data['batch_t']
        batch_r = data['batch_r']
        mode = data['mode']
        h = self.ent_embeddings(batch_h)
        t = self.ent_embeddings(batch_t)
        r = self.rel_embeddings(batch_r)
        # 计算score
        score = self._calc(h,t,r,mode)
        if self.margin_flag:
            return self.margin - score
        else:
            return score
    # 正则化函数
    def regularization(self,data):
        batch_h = data['batch_h']
        batch_t = data['batch_t']
        batch_r = data['batch_r']
        h = self.ent_embeddings(batch_h)
        t = self.ent_embeddings(batch_t)
        r = self.ent_embeddings(batch_r)
        regul = (torch.mean(h ** 2) + torch.mean(t ** 2) + torch.mean(r ** 2)) / 3
        return regul
    def predict(self,data):
        score = self.forward(data)
        if self.margin_flag:
            score = self.margin - score
            return score.cpu().data.numpy()
        else:
            return score.cpu().data.numpy()
    def save_checkpoint(self, path):
        torch.save(self.state_dict(), path)


'''
model = TransE(train_dataloader.get_ent_tot(),train_dataloader.get_rel_tot())
example = list(train_dataloader)[0]
for key in example:
    if type(example[key]) != str:
        example[key] = torch.LongTensor(example[key])
model(example)
'''
pass

3.损失函数构建

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import numpy as np
class MarginLoss(nn.Module):
    def __init__(self,adv_temperature = None,margin = 6.0):
        super(MarginLoss,self).__init__()
        self.margin = nn.Parameter(torch.Tensor([margin]))
        self.margin.requires_grad = False
        if adv_temperature != None:
            self.adv_temperature = nn.Parameter(torch.Tensor([adv_temperature]))
            self.adv_temperatrue.requires_grad = False
            self.adv_flag = True
        else:
            self.adv_flag = False
    def get_weights(self,n_score):
        return F.softmax(-n_score * self.adv_temperature,dim = -1).detach()
    def forward(self,p_score,n_score):
        if self.adv_flag:
            return (self.get_weights(n_score) * torch.max(p_score - n_score,self.margin)).sum(dim = -1).mean() + self.margin
        else:
            return (torch.max(p_score - n_score,-self.margin)).mean() + self.margin
    def predict(self,p_score,n_score):
        score = self.forward(p_score,n_score)
        return score.cpu().data.numpy()
'''
model = TransE(train_dataloader.get_ent_tot(),train_dataloader.get_rel_tot())
example = list(train_dataloader)[0]
for key in example:
    if type(example[key]) != str:
        example[key] = torch.LongTensor(example[key])
loss = MarginLoss()
negativeSampling = NegativeSampling(model,loss,batch_size=train_dataloader.batch_size)
negativeSampling(example)
'''
pass

4.负采样构建

class NegativeSampling(nn.Module):
    def __init__(self,model = None,loss = None,batch_size = 256,regul_rate = 0.0,l3_regul_rate = 0.0):
        super(NegativeSampling,self).__init__()
        self.model = model
        self.loss = loss
        self.batch_size = batch_size
        # 正则化系数
        self.regul_rate = regul_rate
        self.l3_regul_rate = l3_regul_rate
    # batch中样本排序,[正,正,正....(batch_size个),负,负....(batch_size*negative_num个)]
    # 第i个正样本对应负样本[i + batch_size,i + 2 * batch_size...i + (negative_num * batch_size)]
    # 从batch中获取正样本
    def _get_positive_score(self,score):
        positive_score = score[:self.batch_size]
        positive_score = positive_score.view(-1,self.batch_size).permute(1,0)
        return positive_score
    def _get_negative_score(self,score):
        negative_score = score[self.batch_size:]
        # 将[i1,j1,k1,i2,j2,k2,i3,j3,k3]
        # 整理成 [i1,j1,k1
        #        i2,j2,k2
        #        i3,j3,k3] 再转转置即可
        negative_score = negative_score.view(-1,self.batch_size).permute(1,0)
        return negative_score
    def forward(self,data):
        # 调用模型
        score = self.model(data)
        # 正样本
        p_score = self._get_positive_score(score)
        # 负样本
        n_score = self._get_negative_score(score)
        # 得到标量loss
        loss_res = self.loss(p_score,n_score)
        # loss加上正则项
        if self.regul_rate != 0:
            loss_res += self.regul_rate * self.model.regularization(data)
        if self.l3_regul_rate != 0:
            loss_res += self.l3_regul_rate * self.model.l3_regularization()
        return loss_res
    def save_checkpoint(self, path):
        torch.save(self.state_dict(), path)

5.模型训练

import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.optim as optim
import os
import time
import sys
import datetime
import ctypes
import json
import numpy as np
import copy
from tqdm import tqdm
class Traniner(object):
    def __init__(self,model = None,
                 data_loader = None,
                 train_times = 1000,
                 alpha = 0.5,
                 use_gpu = True,
                 opt_method = 'sgd',
                 save_steps = None,
                 checkpoint_dir = None
                ):
        self.work_threads = 8
        self.train_times = train_times
        self.opt_method = opt_method
        self.optimizer = None
        # 避免过拟合
        self.weight_decay = 0
        self.alpha = alpha
        self.model = model
        # class TrainDataLoader:训练集
        self.data_loader = data_loader
        self.use_gpu = use_gpu
        self.save_steps = save_steps
        self.checkpoint_dir = checkpoint_dir
    def to_var(self,x,use_gpu):
        if use_gpu:
            return Variable(torch.from_numpy(x).cuda())
        else:
            return Variable(torch.from_numpy(x))
    def train_one_step(self,data):
        self.optimizer.zero_grad()
        loss = self.model({
            'batch_h':self.to_var(data['batch_h'],self.use_gpu),
            'batch_t':self.to_var(data['batch_t'],self.use_gpu),
            'batch_r':self.to_var(data['batch_r'],self.use_gpu),
            'batch_y':self.to_var(data['batch_y'],self.use_gpu),
            'mode':data['mode']
        })
        loss.backward()
        self.optimizer.step()
        return loss.item()
    def run(self):
        if self.use_gpu:
            self.model.cuda()
        if self.optimizer != None:
            pass
        elif self.opt_method == 'Adagrad' or self.opt_method == 'adagrad':
            self.optimizer = optim.Adagrad(
                self.model.parameters(),
                lr = self.alpha,
                lr_decay = self.lr_decay,
                weight_decay = self.weight_decay
            )
        elif self.opt_method == 'Adadelta' or self.opt_method == 'adadelta':
            self.optimizer = optim.Adadelta(
                self.model.parameters(),
                lr = self.alpha,
                weight_decay = self.weight_decay
            )
        elif self.opt_method == 'Adam' or self.opt_method == 'adam':
            self.optimizer = optim.Adam(
                self.model.parameters(),
                lr = self.alpha,
                weight_decay = self.weight_decay
            )
        else:
            self.optimizer = optim.SGD(
                self.model.parameters(),
                lr = self.alpha,
                weight_decay = self.weight_decay
            )
        print('Finish initializing...')
        # 循环1000个epoch
        training_range = tqdm(range(self.train_times))
        for epoch in training_range:
            res = 0.0
            for data in self.data_loader:
                loss = self.train_one_step(data)
                res += loss
            training_range.set_description('Epoch %d | loss: %f' % (epoch,res))
            if self.save_steps and self.checkpoint_dir and (epoch + 1) % self.save_steps == 0:
                print('Epoch %d has finished,saving...' % (epoch))
                self.model.save_checkpoint(os.path.join(self.checkpoint_dir + '-' + str(epoch) + '.ckpt'))

transe = TransE(train_dataloader.get_ent_tot(),train_dataloader.get_rel_tot())
train_dataloader = TrainDataLoader(
    in_path = "./benchmarks/FB15K237_tiny/",
    nbatches = 100,
    threads = 8,
    # 负采样
    sampling_mode = 'normal',
    # bern构建负样本方式
    bern_flag = 1,
    # 负样本同replace entities
    neg_ent = 25,
    neg_rel = 0
)
# dataloader for test
test_dataloader = TestDataLoader("./benchmarks/FB15K237_tiny/", "link")
loss = MarginLoss()
model = NegativeSampling(transe,loss,batch_size=train_dataloader.batch_size)
trainer = Traniner(model = model,data_loader = train_dataloader)
trainer.run()