def __init__(self, kg, batch_size, sampling_type, use_cuda=None): self.h = kg.head_idx self.t = kg.tail_idx self.r = kg.relations self.m = kg.magnitudes self.use_cuda = use_cuda self.b_size = batch_size if sampling_type == "unif": self.sampler = UniformNegativeSampler(kg) elif sampling_type == "bern": self.sampler = BernoulliNegativeSampler(kg) self.tmp_cuda = use_cuda in ["batch", "all"] if use_cuda is not None and use_cuda == "all": self.h = self.h.cuda() self.t = self.t.cuda() self.r = self.r.cuda() self.m = self.m.cuda()
def main(): # Define some hyper-parameters for training global optimizer benchmarks = 'GeoDBpedia21' model_name = 'TransR_GDR' opt_method = 'Adam' # "Adagrad" "Adadelta" "Adam" "SGD" GDR = True # 是否引入坐标信息 emb_dim = 100 # TransE model ent_dim = emb_dim rel_dim = emb_dim lr = 0.001 margin = 0.5 n_epochs = 20000 train_b_size = 256 # 训练时batch size eval_b_size = 64 # 测评valid test 时batch size # save_time_freq = 5 # require_improvement = save_time_freq*5 validation_freq = 10 # 多少轮进行在验证集进行一次测试 同时保存最佳模型 require_improvement = validation_freq * 3 # 验证集top_k超过多少epoch没下降,结束训练 model_save_path = './checkpoint/' + benchmarks + '_' + model_name + '_' + opt_method + '.ckpt' # 保存最佳hits k (ent)模型 device = 'cuda:0' if cuda.is_available() else 'cpu' # Load dataset module = getattr(import_module('torchkge.models'), model_name + 'Model') load_data = getattr(import_module('torchkge.utils.datasets'), 'load_' + benchmarks) print('Loading data...') kg_train, kg_val, kg_test = load_data(GDR=GDR) print(f'Train set: {kg_train.n_ent} entities, {kg_train.n_rel} relations, {kg_train.n_facts} triplets.') print(f'Valid set: {kg_val.n_facts} triplets, Test set: {kg_test.n_facts} triplets.') # Define the model and criterion print('Loading model...') if 'TransE' in model_name: model = module(emb_dim, kg_train.n_ent, kg_train.n_rel, dissimilarity_type='L2') else: model = module(ent_dim, rel_dim, kg_train.n_ent, kg_train.n_rel) criterion = MarginLoss(margin) # Move everything to CUDA if available if device == 'cuda:0': cuda.empty_cache() model.to(device) criterion.to(device) dataloader = DataLoader(kg_train, batch_size=train_b_size, use_cuda='all') else: dataloader = DataLoader(kg_train, batch_size=train_b_size, use_cuda=None) # Define the torch optimizer to be used optimizer = optimizer(model, opt_method=opt_method, lr=lr) # optimizer = Adam(model.parameters(), lr=lr, weight_decay=1e-5) sampler = BernoulliNegativeSampler(kg_train) start_epoch = 1 best_score = float('-inf') if os.path.exists(model_save_path): # 存在则加载模型 并继续训练 start_epoch, best_score = load_ckpt(model_save_path, model, optimizer) print(f'loading ckpt sucessful, start on epoch {start_epoch}...') print(model) print('lr: {}, margin: {}, dim {}, total epoch: {}, device: {}, batch size: {}, optim: {}, GDR: {}' \ .format(lr, margin, emb_dim, n_epochs, device, train_b_size, opt_method, GDR)) print('Training...') last_improve = start_epoch # 记录上次验证集loss下降的epoch数 start = time.time() # last_improve = start # save_time = start for epoch in range(start_epoch, n_epochs + 1): # model.normalize_parameters() running_loss = 0.0 model.train() for i, batch in enumerate(dataloader): if GDR: h, t, r, point = batch[0], batch[1], batch[2], batch[3] n_h, n_t = sampler.corrupt_batch(h, t, r) # 1:1 negative sampling n_point = id2point(n_h, n_t, kg_train.id2point) optimizer.zero_grad() # forward + backward + optimize pos, neg = model(h, t, n_h, n_t, r) loss = criterion(pos, neg, point, n_point) else: h, t, r = batch[0], batch[1], batch[2] n_h, n_t = sampler.corrupt_batch(h, t, r) optimizer.zero_grad() pos, neg = model(h, t, n_h, n_t, r) loss = criterion(pos, neg) loss.backward() optimizer.step() running_loss += loss.item() model.normalize_parameters() # print('\rEpoch [{:>4}/{:>4}] | mean loss: {:>8.3f}, time: {}'.format(epoch, n_epochs, running_loss / len(dataloader), time_since(start)), end='', flush=True) # # test if epoch % validation_freq == 0: create_dir_not_exists('./checkpoint') model.eval() evaluator = LinkPredictionEvaluator(model, kg_val) evaluator.evaluate(b_size=eval_b_size, verbose=False) _, hit_at_k = evaluator.hit_at_k(10) # val filter hit_k print('Epoch [{:>5}/{:>5}] '.format(epoch, n_epochs), end='') if hit_at_k > best_score: save_ckpt(model, optimizer, epoch, best_score, model_save_path) best_score = hit_at_k improve = '*' # 在有提升的结果后面加上*标注 last_improve = epoch # 验证集hit_k增大即认为有提升 else: improve = '' msg = '| Train loss: {:>8.3f}, Val Hit@10: {:>5.2%}, Time {} {}' print(msg.format(running_loss / len(dataloader), hit_at_k, time_since(start), improve)) # model.normalize_parameters() if epoch - last_improve > require_improvement: # 验证集top_k超过一定epoch没增加,结束训练 print("\nNo optimization for a long time, auto-stopping...") break # # test # if (time.time() - save_time)/60 > save_time_freq: # create_dir_not_exists('./checkpoint') # model.eval() # evaluator = LinkPredictionEvaluator(model, kg_val) # evaluator.evaluate(b_size=eval_b_size, verbose=False) # _, hit_at_k = evaluator.hit_at_k(10) # val filter hit_k # if hit_at_k > best_score: # save_ckpt(model, optimizer, epoch, best_score, model_save_path) # best_score = hit_at_k # improve = '*' # 在有提升的结果后面加上*标注 # last_improve = time.time() # 验证集hit_k增大即认为有提升 # else: # improve = '' # save_time = time.time() # msg = ', Val Hit@10: {:>5.2%} {}' # print(msg.format(hit_at_k, improve)) # model.normalize_parameters() # if (time.time() - last_improve)/60 > require_improvement: # # 验证集top_k超过一定epoch没增加,结束训练 # print("\nNo optimization for a long time, auto-stopping...") # break print('\nTraining done, start evaluate on test data...') print('model name: {}, lr: {}, dim {}, device: {}, eval batch size: {}, optim: {}, GDR: {}' \ .format(model_name, lr, emb_dim, device, eval_b_size, opt_method, GDR)) # Testing the best checkpoint on test dataset load_ckpt(model_save_path, model, optimizer) model.eval() lp_evaluator = LinkPredictionEvaluator(model, kg_test) lp_evaluator.evaluate(eval_b_size, verbose=False) lp_evaluator.print_results() rp_evaluator = RelationPredictionEvaluator(model, kg_test) rp_evaluator.evaluate(eval_b_size, verbose=False) rp_evaluator.print_results() print(f'Total time cost: {time_since(start)}')
criterion = MarginLoss(margin) # In[43]: # Move everything to CUDA if available if cuda.is_available(): cuda.empty_cache() model.cuda() criterion.cuda() # In[20]: # Define the torch optimizer to be used optimizer = Adam(model.parameters(), lr=lr, weight_decay=1e-5) sampler = BernoulliNegativeSampler(kg_train) dataloader = DataLoader(kg_train, batch_size=b_size, use_cuda='all') iterator = tqdm(range(n_epochs), unit='epoch') for epoch in iterator: running_loss = 0.0 for i, batch in enumerate(dataloader): h, t, r = batch[0], batch[1], batch[2] n_h, n_t = sampler.corrupt_batch(h, t, r) optimizer.zero_grad() # forward + backward + optimize pos, neg = model(h, t, n_h, n_t, r) loss = criterion(pos, neg) loss.backward()