def train_data_iterator(entities, triples): entities_1, entities_2 = entities triples_1, triples_2 = triples loader_head_1 = DataLoader(TrainDataset(triples_1, entities_1, config.neg_size, "head-batch"), batch_size=config.batch_size, shuffle=True, num_workers=max(0, config.cpu_num // 3), collate_fn=TrainDataset.collate_fn) loader_tail_1 = DataLoader(TrainDataset(triples_1, entities_1, config.neg_size, "tail-batch"), batch_size=config.batch_size, shuffle=True, num_workers=max(0, config.cpu_num // 3), collate_fn=TrainDataset.collate_fn) loader_head_2 = DataLoader(TrainDataset(triples_2, entities_2, config.neg_size, "head-batch"), batch_size=config.batch_size, shuffle=True, num_workers=max(0, config.cpu_num // 3), collate_fn=TrainDataset.collate_fn) loader_tail_2 = DataLoader(TrainDataset(triples_2, entities_2, config.neg_size, "tail-batch"), batch_size=config.batch_size, shuffle=True, num_workers=max(0, config.cpu_num // 3), collate_fn=TrainDataset.collate_fn) return BidirectionalOneShotIterator(loader_head_1, loader_tail_1, loader_head_2, loader_tail_2)
def train_NoiGAN(trainer): trainer.embedding_model.eval() st = time.time() trainer.positive_triples = trainer.find_positive_triples() et = time.time() print("take %d s to find positive triples" % (et - st)) trainer.train_dataset_head = TrainDataset( trainer.train_triples, trainer.args.nentity, trainer.args.nrelation, trainer.args.negative_sample_size, "head-batch") trainer.train_dataset_head.triples = trainer.positive_triples trainer.train_dataset_tail = TrainDataset( trainer.train_triples, trainer.args.nentity, trainer.args.nrelation, trainer.args.negative_sample_size, "tail-batch") trainer.train_dataset_tail.triples = trainer.positive_triples trainer.train_dataloader_head = DataLoader( trainer.train_dataset_head, batch_size=128, shuffle=True, num_workers=5, collate_fn=TrainDataset.collate_fn) trainer.train_dataloader_tail = DataLoader( trainer.train_dataset_tail, batch_size=128, shuffle=True, num_workers=5, collate_fn=TrainDataset.collate_fn) trainer.train_iterator = BidirectionalOneShotIterator( trainer.train_dataloader_head, trainer.train_dataloader_tail) epochs = 1500 epoch_reward, epoch_loss, avg_reward = 0, 0, 0 for i in range(epochs): trainer.generator.train() positive_sample, negative_sample, subsampling_weight, mode = next( trainer.train_iterator) if trainer.args.cuda: positive_sample = positive_sample.cuda() # [batch_size, 3] negative_sample = negative_sample.cuda( ) # [batch_size, negative_sample_size] #$ embed() pos, neg, scores, sample_idx, row_idx = trainer.generate( positive_sample, negative_sample, mode) loss, rewards = trainer.discriminate(pos, neg, mode) epoch_reward += torch.sum(rewards) epoch_loss += loss rewards = rewards - avg_reward trainer.generator.zero_grad() log_probs = F.log_softmax(scores, dim=1) reinforce_loss = torch.sum( Variable(rewards) * log_probs[row_idx.cuda(), sample_idx.data]) reinforce_loss.backward() trainer.gen_optimizer.step() trainer.generator.eval()
def train_data_iterator(train_triples, ent_num): modes = ["head-batch", "tail-batch"] datasets = [ DataLoader(TrainDataset(train_triples, ent_num, config.neg_size, mode), batch_size=config.batch_size, shuffle=True, num_workers=4, collate_fn=TrainDataset.collate_fn) for mode in modes ] return BidirectionalOneShotIterator(datasets[0], datasets[1])
def train_data_iterator(train_triples, ent_num): dataloader_head = DataLoader(TrainDataset(train_triples, ent_num, config.neg_size, "head-batch"), batch_size=config.batch_size, shuffle=True, num_workers=max(0, config.cpu_num // 3), collate_fn=TrainDataset.collate_fn) dataloader_tail = DataLoader(TrainDataset(train_triples, ent_num, config.neg_size, "tail-batch"), batch_size=config.batch_size, shuffle=True, num_workers=max(0, config.cpu_num // 3), collate_fn=TrainDataset.collate_fn) return BidirectionalOneShotIterator(dataloader_head, dataloader_tail)
def init_dataset(self): train_dataloader_head = DataLoader(TrainDataset( self.train_triples, self.entity_count, self.attr_count, self.value_count, 512, 'head-batch'), batch_size=1024, shuffle=False, num_workers=4, collate_fn=TrainDataset.collate_fn) train_dataloader_tail = DataLoader(TrainDataset( self.train_triples, self.entity_count, self.attr_count, self.value_count, 512, 'tail-batch'), batch_size=1024, shuffle=False, num_workers=4, collate_fn=TrainDataset.collate_fn) self.train_iterator = BidirectionalOneShotIterator( train_dataloader_head, train_dataloader_tail)
def construct_dataloader(args, train_triples, nentity, nrelation): train_dataloader_head = DataLoader(TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, 'head-batch'), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn) train_dataloader_tail = DataLoader(TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, 'tail-batch'), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn) train_iterator = BidirectionalOneShotIterator(train_dataloader_head, train_dataloader_tail) return train_iterator
def main(args): # if (not args.do_train) and (not args.do_valid) and (not args.do_test) and (not args.do_case) and (not args.fire_test) and (not args.rel_do_test) : # raise ValueError('one of train/val/test mode must be choosed.') if args.init_checkpoint: override_config(args) elif args.data_path is None: raise ValueError('one of init_checkpoint/data_path must be choosed.') if args.do_train and args.save_path is None: raise ValueError('Where do you want to save your trained model?') if args.save_path and not os.path.exists(args.save_path): os.makedirs(args.save_path) # Write logs to checkpoint and console set_logger(args) with open(os.path.join(args.data_path, 'entities.dict')) as fin: entity2id = dict() for line in fin: eid, entity = line.strip().split('\t') entity2id[entity] = int(eid) with open(os.path.join(args.data_path, 'relations.dict')) as fin: relation2id = dict() for line in fin: rid, relation = line.strip().split('\t') relation2id[relation] = int(rid) # Read regions for Countries S* datasets if args.countries: regions = list() with open(os.path.join(args.data_path, 'regions.list')) as fin: for line in fin: region = line.strip() regions.append(entity2id[region]) args.regions = regions nentity = len(entity2id) nrelation = len(relation2id) args.nentity = nentity args.nrelation = nrelation logging.info('Model: %s' % args.model) logging.info('Data Path: %s' % args.data_path) logging.info('#entity: %d' % nentity) logging.info('#relation: %d' % nrelation) train_triples = read_triple(os.path.join(args.data_path, 'train.txt'), entity2id, relation2id) # train_triples = read_triple(os.path.join(args.data_path, 'train_1900.txt'), entity2id, relation2id) logging.info('#train: %d' % len(train_triples)) valid_triples = read_triple(os.path.join(args.data_path, 'valid.txt'), entity2id, relation2id) logging.info('#valid: %d' % len(valid_triples)) test_triples = read_triple(os.path.join(args.data_path, 'test.txt'), entity2id, relation2id) # seen_test_triples = read_triple(os.path.join(args.data_path, 'seen_test.txt'), entity2id, relation2id) # test_triples = read_triple(os.path.join(args.data_path, 'test_alone_triples.txt'), entity2id, relation2id) # def file_name(file_dir): # for root, dirs, files in os.walk(file_dir): # return files # rel_dataset = file_name("/scratch/mengyali/workspace/rotate/data/wn18rr/rel_dataset_txt/") # for rel in rel_dataset: # test_triples = read_triple(os.path.join(args.data_path, "rel_dataset_txt/"+str(rel)), entity2id, relation2id) # logging.info('#test: %d' % len(test_triples)) #All true triples all_true_triples = train_triples + valid_triples + test_triples kge_model = KGEModel( model_name=args.model, nentity=nentity, nrelation=nrelation, hidden_dim=args.hidden_dim, gamma=args.gamma, double_entity_embedding=args.double_entity_embedding, double_relation_embedding=args.double_relation_embedding) logging.info('Model Parameter Configuration:') for name, param in kge_model.named_parameters(): logging.info('Parameter %s: %s, require_grad = %s' % (name, str(param.size()), str(param.requires_grad))) if args.cuda: kge_model = kge_model.cuda() if args.do_train: # Set training dataloader iterator train_dataloader_head = DataLoader( TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, 'head-batch', entity2id, relation2id, args.data_path, args.typecons), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn) train_dataloader_tail = DataLoader( TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, 'tail-batch', entity2id, relation2id, args.data_path, args.typecons), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn) train_iterator = BidirectionalOneShotIterator(train_dataloader_head, train_dataloader_tail) # Set training configuration current_learning_rate = args.learning_rate optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, kge_model.parameters()), lr=current_learning_rate) if args.warm_up_steps: warm_up_steps = args.warm_up_steps else: warm_up_steps = args.max_steps // 2 if args.init_checkpoint: # Restore model from checkpoint directory logging.info('Loading checkpoint %s...' % args.init_checkpoint) checkpoint = torch.load( os.path.join(args.init_checkpoint, 'checkpoint')) init_step = checkpoint['step'] kge_model.load_state_dict(checkpoint['model_state_dict']) if args.do_train: current_learning_rate = checkpoint['current_learning_rate'] warm_up_steps = checkpoint['warm_up_steps'] optimizer.load_state_dict(checkpoint['optimizer_state_dict']) else: logging.info('Ramdomly Initializing %s Model...' % args.model) init_step = 0 step = init_step logging.info('Start Training...') logging.info('init_step = %d' % init_step) logging.info('batch_size = %d' % args.batch_size) logging.info('negative_adversarial_sampling = %d' % args.negative_adversarial_sampling) logging.info('hidden_dim = %d' % args.hidden_dim) logging.info('gamma = %f' % args.gamma) logging.info('negative_adversarial_sampling = %s' % str(args.negative_adversarial_sampling)) if args.negative_adversarial_sampling: logging.info('adversarial_temperature = %f' % args.adversarial_temperature) # Set valid dataloader as it would be evaluated during training if args.do_train: logging.info('learning_rate = %d' % current_learning_rate) training_logs = [] #Training Loop for step in range(init_step, args.max_steps): log = kge_model.train_step(kge_model, optimizer, train_iterator, args) training_logs.append(log) if step >= warm_up_steps: current_learning_rate = current_learning_rate / 10 logging.info('Change learning_rate to %f at step %d' % (current_learning_rate, step)) optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, kge_model.parameters()), lr=current_learning_rate) warm_up_steps = warm_up_steps * 3 if step % args.save_checkpoint_steps == 0: save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, 'warm_up_steps': warm_up_steps } save_model(kge_model, optimizer, save_variable_list, args) if step % args.log_steps == 0: metrics = {} for metric in training_logs[0].keys(): metrics[metric] = sum( [log[metric] for log in training_logs]) / len(training_logs) log_metrics('Training average', step, metrics) training_logs = [] if args.do_valid and step % args.valid_steps == 0: logging.info('Evaluating on Valid Dataset...') metrics = kge_model.test_step(kge_model, valid_triples, all_true_triples, args) log_metrics('Valid', step, metrics) save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, 'warm_up_steps': warm_up_steps } save_model(kge_model, optimizer, save_variable_list, args) if args.do_valid: logging.info('Evaluating on Valid Dataset...') metrics = kge_model.test_step(kge_model, valid_triples, all_true_triples, args) log_metrics('Valid', step, metrics) if args.do_test: logging.info('Evaluating on Test Dataset...') metrics = kge_model.test_step(kge_model, test_triples, all_true_triples, args) # logging.info("----------------------"+str(rel)+"---------------------\n") log_metrics('Test', step, metrics) if args.get_metric: logging.info( 'Evaluating on Test Dataset and Show the metric in two sides...') head_metrics, tail_metrics = kge_model.get_metric( kge_model, test_triples, all_true_triples, args) logging.info("--------------- Head ------------\n") log_metrics('Test-Head', step, head_metrics) logging.info("--------------- Tail ------------\n") log_metrics('Test-Tail', step, tail_metrics) if args.evaluate_train: logging.info('Evaluating on Training Dataset...') metrics = kge_model.test_step(kge_model, train_triples, all_true_triples, args) log_metrics('Test', step, metrics) # Codes about StAR if args.get_scores: for type in ['dev', 'test']: kge_model.get_scores(kge_model, type, valid_triples, all_true_triples, args) if args.get_model_dataset: kge_model.get_model_dataset(kge_model, 'train', valid_triples, all_true_triples, args) if args.get_cases: kge_model.get_cases(kge_model, test_triples, all_true_triples, args) if args.rel_do_test: train_rel_dict = collections.Counter([ex[1] for ex in train_triples]) rel_dict = dict() logging.info('Evaluating on Each Test Dataset Devided by Relation...') test_ex_dict_departby_rel = dict.fromkeys(relation2id.keys(), []) for _rel in relation2id: test_ex_dict_departby_rel[_rel] = [ _ex for _ex in test_triples if _ex[1] == relation2id[_rel] ] for _rel in test_ex_dict_departby_rel.keys(): _rel_test_triples = test_ex_dict_departby_rel[_rel] _rel_data = [ train_rel_dict[relation2id[_rel]], len(_rel_test_triples) ] if len(_rel_test_triples) != 0: metrics = kge_model.test_step(kge_model, _rel_test_triples, all_true_triples, args) _rel_data.extend([ round(metrics['HITS@1'], 3), round(metrics['HITS@3'], 3), round(metrics['HITS@10'], 3), round(metrics['MR'], 1), round(metrics['MRR'], 3) ]) else: _rel_data.extend([0, 0, 0, 0, 0]) rel_dict[_rel] = _rel_data sorted_rel = sorted(rel_dict.items(), key=lambda x: x[1][0], reverse=True) save_dir = args.init_checkpoint with open(join(save_dir, "rel_unbalanced.txt"), "w", encoding="utf-8") as fp: fp.write(str(sorted_rel)) torch.save(sorted_rel, join(save_dir, "rel_unbalanced")) # SaveInExcle(sorted_rel, save_dir) print("explore unbalanced finished")
def main(args): if args.seed != -1: torch.manual_seed(args.seed) if args.cuda: torch.cuda.manual_seed(args.seed) if (not args.do_train) and (not args.do_valid) and (not args.do_test): raise ValueError('one of train/val/test mode must be choosed.') if args.init_checkpoint: override_config(args) elif args.data_path is None: raise ValueError('one of init_checkpoint/data_path must be choosed.') if args.do_train and args.save_path is None: raise ValueError('Where do you want to save your trained model?') if args.save_path and not os.path.exists(args.save_path): os.makedirs(args.save_path) if args.do_train and args.do_valid: if not os.path.exists("%s/best/" % args.save_path): os.makedirs("%s/best/" % args.save_path) # Write logs to checkpoint and console set_logger(args) with open(os.path.join(args.data_path, 'entities.dict')) as fin: entity2id = dict() for line in fin: eid, entity = line.strip().split('\t') entity2id[entity] = int(eid) with open(os.path.join(args.data_path, 'relations.dict')) as fin: relation2id = dict() for line in fin: rid, relation = line.strip().split('\t') relation2id[relation] = int(rid) # Read regions for Countries S* datasets if args.countries: regions = list() with open(os.path.join(args.data_path, 'regions.list')) as fin: for line in fin: region = line.strip() regions.append(entity2id[region]) args.regions = regions nentity = len(entity2id) nrelation = len(relation2id) args.nentity = nentity args.nrelation = nrelation logging.info('Model: %s' % args.model) logging.info('Data Path: %s' % args.data_path) logging.info('#entity: %d' % nentity) logging.info('#relation: %d' % nrelation) train_triples = read_triple(os.path.join(args.data_path, 'train.txt'), entity2id, relation2id) logging.info('#train: %d' % len(train_triples)) valid_triples = read_triple(os.path.join(args.data_path, 'valid.txt'), entity2id, relation2id) logging.info('#valid: %d' % len(valid_triples)) test_triples = read_triple(os.path.join(args.data_path, 'test.txt'), entity2id, relation2id) logging.info('#test: %d' % len(test_triples)) train_triples_tsr = torch.LongTensor(train_triples).transpose( 0, 1) #idx X batch #All true triples all_true_triples = train_triples + valid_triples + test_triples #if args.use_gnn: # assert False # #kge_model = GNN_KGEModel( # # model_name=args.model, # # nentity=nentity, # # nrelation=nrelation, # # hidden_dim=args.hidden_dim, # # gamma=args.gamma, # # num_layers=args.gnn_layers, # # args = args, # # dropout=args.dropout, # # double_entity_embedding=args.double_entity_embedding, # # double_relation_embedding=args.double_relation_embedding, # #) #else: kge_model = KGEModel( model_name=args.model, nentity=nentity, nrelation=nrelation, hidden_dim=args.hidden_dim, gamma=args.gamma, args=args, double_entity_embedding=args.double_entity_embedding, double_relation_embedding=args.double_relation_embedding, ) logging.info('Model Configuration:') logging.info(str(kge_model)) logging.info('Model Parameter Configuration:') for name, param in kge_model.named_parameters(): logging.info('Parameter %s: %s, require_grad = %s' % (name, str(param.size()), str(param.requires_grad))) if args.cuda: kge_model = kge_model.cuda() train_triples_tsr = train_triples_tsr.cuda() #kge_model.build_cxt_triple_map(train_triples) if args.do_train: # Set training dataloader iterator if args.same_head_tail: #shuffle train_triples first and no shuffle within dataloaders. So both head and tail will share the same idx shuffle(train_triples) train_dataloader_head = DataLoader( TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, 'head-batch'), batch_size=args.batch_size, shuffle=False, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn) train_dataloader_tail = DataLoader( TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, 'tail-batch'), batch_size=args.batch_size, shuffle=False, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn) else: train_dataloader_head = DataLoader( TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, 'head-batch'), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn) train_dataloader_tail = DataLoader( TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, 'tail-batch'), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn) train_iterator = BidirectionalOneShotIterator(train_dataloader_head, train_dataloader_tail) #else: # train_dataloader_rel = DataLoader( # TrainDataset(train_triples, nentity, nrelation, # args.negative_sample_head_size*args.negative_sample_tail_size, # 'rel-batch', # negative_sample_head_size =args.negative_sample_head_size, # negative_sample_tail_size =args.negative_sample_tail_size, # half_correct=args.negative_sample_half_correct), # batch_size=args.batch_size, # shuffle=True, # num_workers=max(1, args.cpu_num//2), # collate_fn=TrainDataset.collate_fn # ) # train_iterator = BidirectionalOneShotIterator.one_shot_iterator(train_dataloader_rel) # tail_only = True # Set training configuration current_learning_rate = args.learning_rate optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, kge_model.parameters()), lr=current_learning_rate, weight_decay=args.weight_decay, ) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.5, last_epoch=-1) #if args.warm_up_steps: # warm_up_steps = args.warm_up_steps #else: # warm_up_steps = args.max_steps // 2 if args.init_checkpoint: # Restore model from checkpoint directory logging.info('Loading checkpoint %s...' % args.init_checkpoint) checkpoint = torch.load( os.path.join(args.init_checkpoint, 'checkpoint')) init_step = checkpoint['step'] if 'score_weight' in kge_model.state_dict( ) and 'score_weight' not in checkpoint['model_state_dict']: checkpoint['model_state_dict'][ 'score_weights'] = kge_model.state_dict()['score_weights'] kge_model.load_state_dict(checkpoint['model_state_dict']) if args.do_train: current_learning_rate = checkpoint['current_learning_rate'] #warm_up_steps = checkpoint['warm_up_steps'] optimizer.load_state_dict(checkpoint['optimizer_state_dict']) else: current_learning_rate = 0 elif args.init_embedding: logging.info('Loading pretrained embedding %s ...' % args.init_embedding) if kge_model.entity_embedding is not None: entity_embedding = np.load( os.path.join(args.init_embedding, 'entity_embedding.npy')) relation_embedding = np.load( os.path.join(args.init_embedding, 'relation_embedding.npy')) entity_embedding = torch.from_numpy(entity_embedding).to( kge_model.entity_embedding.device) relation_embedding = torch.from_numpy(relation_embedding).to( kge_model.relation_embedding.device) kge_model.entity_embedding.data[:entity_embedding. size(0)] = entity_embedding kge_model.relation_embedding.data[:relation_embedding. size(0)] = relation_embedding init_step = 1 current_learning_rate = 0 else: logging.info('Ramdomly Initializing %s Model...' % args.model) init_step = 1 step = init_step logging.info('Start Training...') logging.info('init_step = %d' % init_step) logging.info('learning_rate = %.5f' % current_learning_rate) logging.info('batch_size = %d' % args.batch_size) logging.info('negative_adversarial_sampling = %d' % args.negative_adversarial_sampling) logging.info('hidden_dim = %d' % args.hidden_dim) logging.info('gamma = %f' % args.gamma) logging.info('negative_adversarial_sampling = %s' % str(args.negative_adversarial_sampling)) if args.negative_adversarial_sampling: logging.info('adversarial_temperature = %f' % args.adversarial_temperature) # Set valid dataloader as it would be evaluated during training #loss_func = nn.BCEWithLogitsLoss(reduction="none") if args.use_bceloss else nn.LogSigmoid() if args.use_bceloss: loss_func = nn.BCELoss(reduction="none") elif args.use_softmarginloss: loss_func = nn.SoftMarginLoss(reduction="none") else: loss_func = nn.LogSigmoid() #kge_model.cluster_relation_entity_embedding(args.context_cluster_num, args.context_cluster_scale) if args.do_train: training_logs = [] best_metrics = None #Training Loop optimizer.zero_grad() for step in range(init_step, args.max_steps + 1): if step % args.update_freq == 1 or args.update_freq == 1: optimizer.zero_grad() log = kge_model.train_step(kge_model, train_iterator, train_triples_tsr, loss_func, args) if step % args.update_freq == 0: optimizer.step() training_logs.append(log) #if step >= warm_up_steps: # current_learning_rate = current_learning_rate / 10 # logging.info('Change learning_rate to %f at step %d' % (current_learning_rate, step)) # optimizer = torch.optim.Adam( # filter(lambda p: p.requires_grad, kge_model.parameters()), # lr=current_learning_rate # ) # warm_up_steps = warm_up_steps * 3 if step % args.schedule_steps == 0: scheduler.step() if step % args.save_checkpoint_steps == 0: save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, #'warm_up_steps': warm_up_steps } save_model(kge_model, optimizer, save_variable_list, args) if step % args.log_steps == 0: metrics = {} for metric in training_logs[0].keys(): metrics[metric] = sum( [log[metric] for log in training_logs]) / len(training_logs) log_metrics('Training average', step, [metrics]) training_logs = [] if args.do_valid and step % args.valid_steps == 0: logging.info('Evaluating on Valid Dataset...') metrics = kge_model.test_step(kge_model, valid_triples, all_true_triples, train_triples_tsr, args) log_metrics('Valid', step, metrics) if is_better_metric(best_metrics, metrics): save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, #'warm_up_steps': warm_up_steps } save_model(kge_model, optimizer, save_variable_list, args, True) best_metrics = metrics #kge_model.cluster_relation_entity_embedding(args.context_cluster_num, args.context_cluster_scale) save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, #'warm_up_steps': warm_up_steps } save_model(kge_model, optimizer, save_variable_list, args) if args.do_valid and args.do_train: #load the best model best_checkpoint = torch.load("%s/best/checkpoint" % args.save_path) kge_model.load_state_dict(best_checkpoint['model_state_dict']) logging.info("Loading best model from step %d" % best_checkpoint['step']) step = best_checkpoint['step'] if args.do_valid: logging.info('Evaluating on Valid Dataset...') metrics = kge_model.test_step(kge_model, valid_triples, all_true_triples, train_triples_tsr, args) log_metrics('Valid', step, metrics) if args.do_test: logging.info('Evaluating on Test Dataset...') metrics = kge_model.test_step(kge_model, test_triples, all_true_triples, train_triples_tsr, args) log_metrics('Test', step, metrics) if args.evaluate_train: logging.info('Evaluating on Training Dataset...') metrics = kge_model.test_step(kge_model, train_triples, all_true_triples, train_triples_tsr, args) log_metrics('Test', step, metrics)
def main(args): if (not args.do_train) and (not args.do_valid) and (not args.do_test): raise ValueError('one of train/val/test mode must be choosed.') if args.init_checkpoint: override_config(args) elif args.data_path is None: raise ValueError('one of init_checkpoint/data_path must be choosed.') if args.do_train and args.save_path is None: raise ValueError('Where do you want to save your trained model?') if args.save_path and not os.path.exists(args.save_path): os.makedirs(args.save_path) # Write logs to checkpoint and console set_logger(args) with open(os.path.join(args.data_path, 'entities.dict')) as fin: entity2id = dict() for line in fin: eid, entity = line.strip().split('\t') entity2id[entity] = int(eid) with open(os.path.join(args.data_path, 'relations.dict')) as fin: relation2id = dict() for line in fin: rid, relation = line.strip().split('\t') relation2id[relation] = int(rid) # Read regions for Countries S* datasets if args.countries: regions = list() with open(os.path.join(args.data_path, 'regions.list')) as fin: for line in fin: region = line.strip() regions.append(entity2id[region]) args.regions = regions nentity = len(entity2id) nrelation = len(relation2id) args.nentity = nentity args.nrelation = nrelation logging.info('Model: %s' % args.model) logging.info('Data Path: %s' % args.data_path) logging.info('#entity: %d' % nentity) logging.info('#relation: %d' % nrelation) train_triples = read_triple(os.path.join(args.data_path, 'train.txt'), entity2id, relation2id) logging.info('#train: %d' % len(train_triples)) valid_triples = read_triple(os.path.join(args.data_path, 'valid.txt'), entity2id, relation2id) logging.info('#valid: %d' % len(valid_triples)) test_triples = read_triple(os.path.join(args.data_path, 'test.txt'), entity2id, relation2id) logging.info('#test: %d' % len(test_triples)) #All true triples all_true_triples = train_triples + valid_triples + test_triples kge_model = KGEModel( model_name=args.model, nentity=nentity, nrelation=nrelation, hidden_dim=args.hidden_dim, type_dim = args.type_dim, gamma=args.gamma, gamma_type=args.gamma_type, gamma_pair=args.gamma_pair, double_entity_embedding=args.double_entity_embedding, double_relation_embedding=args.double_relation_embedding ) logging.info('Model Parameter Configuration:') for name, param in kge_model.named_parameters(): logging.info('Parameter %s: %s, require_grad = %s' % (name, str(param.size()), str(param.requires_grad))) if args.cuda: kge_model = kge_model.cuda() if args.do_train: # Set training dataloader iterator train_dataloader_head = DataLoader( TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, args.pair_sample_size, 'head-batch'), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num//2), collate_fn=TrainDataset.collate_fn ) train_dataloader_tail = DataLoader( TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, args.pair_sample_size, 'tail-batch'), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num//2), collate_fn=TrainDataset.collate_fn ) train_iterator = BidirectionalOneShotIterator(train_dataloader_head, train_dataloader_tail) # Set training configuration current_learning_rate = args.learning_rate optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, kge_model.parameters()), lr=current_learning_rate ) if args.warm_up_steps: warm_up_steps = args.warm_up_steps else: warm_up_steps = args.max_steps // 2 if args.init_checkpoint: # Restore model from checkpoint directory logging.info('Loading checkpoint %s...' % args.init_checkpoint) checkpoint = torch.load(os.path.join(args.init_checkpoint, 'checkpoint')) init_step = checkpoint['step'] kge_model.load_state_dict(checkpoint['model_state_dict']) if args.do_train: current_learning_rate = checkpoint['current_learning_rate'] warm_up_steps = checkpoint['warm_up_steps'] optimizer.load_state_dict(checkpoint['optimizer_state_dict']) else: logging.info('Ramdomly Initializing %s Model...' % args.model) init_step = 0 step = init_step logging.info('Start Training...') logging.info('init_step = %d' % init_step) logging.info('batch_size = %d' % args.batch_size) logging.info('negative_adversarial_sampling = %d' % args.negative_adversarial_sampling) logging.info('hidden_dim = %d' % args.hidden_dim) logging.info('gamma = %f' % args.gamma) logging.info('type_dim = %d' % args.type_dim) logging.info('gamma_type = %f' % args.gamma_type) logging.info('alpha_1 = %f' % args.alpha_1) logging.info('gamma_pair = %f' % args.gamma_pair) logging.info('alpha_2 = %f' % args.alpha_2) logging.info('negative_adversarial_sampling = %s' % str(args.negative_adversarial_sampling)) logging.info('pair_sample_size = %d' % args.pair_sample_size) if args.negative_adversarial_sampling: logging.info('adversarial_temperature = %f' % args.adversarial_temperature) # Set valid dataloader as it would be evaluated during training if args.do_train: logging.info('learning_rate = %d' % current_learning_rate) training_logs = [] #Training Loop for step in range(init_step, args.max_steps): log = kge_model.train_step(kge_model, optimizer, train_iterator, args) training_logs.append(log) if step >= warm_up_steps: current_learning_rate = current_learning_rate / 10 logging.info('Change learning_rate to %f at step %d' % (current_learning_rate, step)) optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, kge_model.parameters()), lr=current_learning_rate ) warm_up_steps = warm_up_steps * 3 if step % args.save_checkpoint_steps == 0: save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, 'warm_up_steps': warm_up_steps } save_model(kge_model, optimizer, save_variable_list, args) if step % args.log_steps == 0: metrics = {} for metric in training_logs[0].keys(): metrics[metric] = sum([log[metric] for log in training_logs])/len(training_logs) log_metrics('Training average', step, metrics) training_logs = [] if args.do_valid and step % args.valid_steps == 0: logging.info('Evaluating on Valid Dataset...') metrics = kge_model.test_step(kge_model, valid_triples, all_true_triples, args) log_metrics('Valid', step, metrics) save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, 'warm_up_steps': warm_up_steps } save_model(kge_model, optimizer, save_variable_list, args) if args.do_valid: logging.info('Evaluating on Valid Dataset...') metrics = kge_model.test_step(kge_model, valid_triples, all_true_triples, args) log_metrics('Valid', step, metrics) if args.do_test: logging.info('Evaluating on Test Dataset...') metrics = kge_model.test_step(kge_model, test_triples, all_true_triples, args) log_metrics('Test', step, metrics) if args.evaluate_train: logging.info('Evaluating on Training Dataset...') metrics = kge_model.test_step(kge_model, train_triples, all_true_triples, args) log_metrics('Test', step, metrics)
def main(args): # if (not args.do_train) and (not args.do_valid) and (not args.do_test): # raise ValueError('one of train/val/test mode must be choosed.') if args.init_checkpoint: override_config(args) elif args.data_path is None: raise ValueError('one of init_checkpoint/data_path must be choosed.') if args.do_train and args.save_path is None: raise ValueError('Where do you want to save your trained model?') if args.save_path and not os.path.exists(args.save_path): os.makedirs(args.save_path) # Write logs to checkpoint and console set_logger(args) # with open(os.path.join(args.data_path, 'entities.dict')) as fin: # entity2id = dict() # id2entity = dict() # for line in fin: # eid, entity = line.strip().split('\t') # entity2id[entity] = int(eid) # id2entity[int(eid)] = entity # with open(os.path.join(args.data_path, 'relations.dict')) as fin: # relation2id = dict() # id2relation = dict() # for line in fin: # rid, relation = line.strip().split('\t') # relation2id[relation] = int(rid) # id2relation[int(rid)] = relation # # Read regions for Countries S* datasets # if args.countries: # regions = list() # with open(os.path.join(args.data_path, 'regions.list')) as fin: # for line in fin: # region = line.strip() # regions.append(entity2id[region]) # args.regions = regions '''amazon dataset''' with open(os.path.join(args.data_path, 'entity2id.txt')) as fin: entity2id = dict() id2entity = dict() for line in fin: if len(line.strip().split('\t')) < 2: continue entity, eid = line.strip().split('\t') entity2id[entity] = int(eid) id2entity[int(eid)] = entity with open(os.path.join(args.data_path, 'relation2id.txt')) as fin: relation2id = dict() id2relation = dict() for line in fin: if len(line.strip().split('\t')) < 2: continue relation, rid = line.strip().split('\t') relation2id[relation] = int(rid) id2relation[int(rid)] = relation nentity = len(entity2id) nrelation = len(relation2id) args.nentity = nentity args.nrelation = nrelation logging.info('Model: %s' % args.model) logging.info('Data Path: %s' % args.data_path) logging.info('#entity: %d' % nentity) logging.info('#relation: %d' % nrelation) # -------------------------------------------------- # Comments by Meng: # During training, pLogicNet will augment the training triplets, # so here we load both the augmented triplets (train.txt) for training and # the original triplets (train_kge.txt) for evaluation. # Also, the hidden triplets (hidden.txt) are also loaded for annotation. # -------------------------------------------------- # train_triples = read_triple(os.path.join(args.workspace_path, 'train_kge.txt'), entity2id, relation2id) # logging.info('#train: %d' % len(train_triples)) # train_original_triples = read_triple(os.path.join(args.data_path, 'train.txt'), entity2id, relation2id) # logging.info('#train original: %d' % len(train_original_triples)) # valid_triples = read_triple(os.path.join(args.data_path, 'valid.txt'), entity2id, relation2id) # logging.info('#valid: %d' % len(valid_triples)) # test_triples = read_triple(os.path.join(args.data_path, 'test.txt'), entity2id, relation2id) # logging.info('#test: %d' % len(test_triples)) # hidden_triples = read_triple(os.path.join(args.workspace_path, 'hidden.txt'), entity2id, relation2id) # logging.info('#hidden: %d' % len(hidden_triples)) train_triples = read_triple(os.path.join(args.workspace_path, 'train_kge.txt'), entity2id, relation2id) logging.info('#train: %d' % len(train_triples)) train_original_triples = read_triple(os.path.join(args.data_path, 'train.txt'), entity2id, relation2id) logging.info('#train original: %d' % len(train_original_triples)) valid_triples = read_triple(os.path.join(args.data_path, 'kg_val_triples_Cell_Phones_and_Accessories.txt'), entity2id, relation2id) logging.info('#valid: %d' % len(valid_triples)) test_triples = read_triple(os.path.join(args.data_path, 'kg_test_triples_Cell_Phones_and_Accessories.txt'), entity2id, relation2id) logging.info('#test: %d' % len(test_triples)) test_candidates = np.load(os.path.join(args.data_path, 'rec_test_candidate100.npz'))['candidates'][:, 1:] # test_candidates = np.load('/common/users/yz956/kg/code/OpenDialKG/cand.npy') # hidden_triples = read_triple(os.path.join(args.workspace_path, 'hidden.txt'), entity2id, relation2id) hidden_triples = read_triple("/common/users/yz956/kg/code/KBRD/data/cpa/cpa/hidden_50.txt", entity2id, relation2id) logging.info('#hidden: %d' % len(hidden_triples)) #All true triples all_true_triples = train_original_triples + valid_triples + test_triples kge_model = KGEModel( model_name=args.model, nentity=nentity, nrelation=nrelation, hidden_dim=args.hidden_dim, gamma=args.gamma, double_entity_embedding=args.double_entity_embedding, double_relation_embedding=args.double_relation_embedding ) logging.info('Model Parameter Configuration:') for name, param in kge_model.named_parameters(): logging.info('Parameter %s: %s, require_grad = %s' % (name, str(param.size()), str(param.requires_grad))) if args.cuda: kge_model = kge_model.cuda() if args.do_train: # Set training dataloader iterator train_dataloader_head = DataLoader( TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, 'head-batch'), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num//2), collate_fn=TrainDataset.collate_fn ) train_dataloader_tail = DataLoader( TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, 'tail-batch'), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num//2), collate_fn=TrainDataset.collate_fn ) train_iterator = BidirectionalOneShotIterator(train_dataloader_head, train_dataloader_tail) # Set training configuration current_learning_rate = args.learning_rate optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, kge_model.parameters()), lr=current_learning_rate ) if args.warm_up_steps: warm_up_steps = args.warm_up_steps else: warm_up_steps = args.max_steps // 2 if args.init_checkpoint: # Restore model from checkpoint directory logging.info('Loading checkpoint %s...' % args.init_checkpoint) checkpoint = torch.load(os.path.join(args.init_checkpoint, 'checkpoint')) init_step = checkpoint['step'] kge_model.load_state_dict(checkpoint['model_state_dict']) if args.do_train: current_learning_rate = checkpoint['current_learning_rate'] warm_up_steps = checkpoint['warm_up_steps'] optimizer.load_state_dict(checkpoint['optimizer_state_dict']) else: logging.info('Ramdomly Initializing %s Model...' % args.model) init_step = 0 step = init_step logging.info('Start Training...') logging.info('init_step = %d' % init_step) logging.info('learning_rate = %d' % current_learning_rate) logging.info('batch_size = %d' % args.batch_size) logging.info('negative_adversarial_sampling = %d' % args.negative_adversarial_sampling) logging.info('hidden_dim = %d' % args.hidden_dim) logging.info('gamma = %f' % args.gamma) logging.info('negative_adversarial_sampling = %s' % str(args.negative_adversarial_sampling)) if args.negative_adversarial_sampling: logging.info('adversarial_temperature = %f' % args.adversarial_temperature) if args.record: local_path = args.workspace_path ensure_dir(local_path) opt = vars(args) with open(local_path + '/opt.txt', 'w') as fo: for key, val in opt.items(): fo.write('{} {}\n'.format(key, val)) # Set valid dataloader as it would be evaluated during training if args.do_train: training_logs = [] #Training Loop for step in range(init_step, args.max_steps): log = kge_model.train_step(kge_model, optimizer, train_iterator, args) training_logs.append(log) if step >= warm_up_steps: current_learning_rate = current_learning_rate / 10 logging.info('Change learning_rate to %f at step %d' % (current_learning_rate, step)) optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, kge_model.parameters()), lr=current_learning_rate ) warm_up_steps = warm_up_steps * 3 if step % args.save_checkpoint_steps == 0: save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, 'warm_up_steps': warm_up_steps } save_model(kge_model, optimizer, save_variable_list, args) if step % args.log_steps == 0: metrics = {} for metric in training_logs[0].keys(): metrics[metric] = sum([log[metric] for log in training_logs])/len(training_logs) log_metrics('Training average', step, metrics) training_logs = [] if args.do_valid and (step + 1) % args.valid_steps == 0: logging.info('Evaluating on Valid Dataset...') metrics, preds = kge_model.test_step(kge_model, valid_triples, all_true_triples, args) log_metrics('Valid', step, metrics) save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, 'warm_up_steps': warm_up_steps } save_model(kge_model, optimizer, save_variable_list, args) if args.do_valid: logging.info('Evaluating on Valid Dataset...') metrics, preds = kge_model.test_step(kge_model, valid_triples, all_true_triples, args) log_metrics('Valid', step, metrics) # -------------------------------------------------- # Comments by Meng: # Save the prediction results of KGE on validation set. # -------------------------------------------------- if args.record: # Save the final results with open(local_path + '/result_kge_valid.txt', 'w') as fo: for metric in metrics: fo.write('{} : {}\n'.format(metric, metrics[metric])) # Save the predictions on test data with open(local_path + '/pred_kge_valid.txt', 'w') as fo: for h, r, t, f, rk, l in preds: fo.write('{}\t{}\t{}\t{}\t{}\n'.format(id2entity[h], id2relation[r], id2entity[t], f, rk)) for e, val in l: fo.write('{}:{:.4f} '.format(id2entity[e], val)) fo.write('\n') if args.do_test: logging.info('Evaluating on Test Dataset...') # metrics, preds = kge_model.test_step(kge_model, test_triples, all_true_triples, args) metrics, preds = kge_model.test_step(kge_model, test_triples, test_candidates, all_true_triples, args) log_metrics('Test', step, metrics) # -------------------------------------------------- # Comments by Meng: # Save the prediction results of KGE on test set. # -------------------------------------------------- if args.record: # Save the final results with open(local_path + '/result_kge.txt', 'w') as fo: for metric in metrics: fo.write('{} : {}\n'.format(metric, metrics[metric])) # Save the predictions on test data with open(local_path + '/pred_kge.txt', 'w') as fo: for h, r, t, f, rk, l in preds: fo.write('{}\t{}\t{}\t{}\t{}\n'.format(id2entity[h], id2relation[r], id2entity[t], f, rk)) for e, val in l: fo.write('{}:{:.4f} '.format(id2entity[e], val)) fo.write('\n') # -------------------------------------------------- # Comments by Meng: # Save the annotations on hidden triplets. # -------------------------------------------------- if args.record: # Annotate hidden triplets scores = kge_model.infer_step(kge_model, hidden_triples, args) # with open(local_path + '/annotation.txt', 'w') as fo: # for (h, r, t), s in zip(hidden_triples, scores): # fo.write('{}\t{}\t{}\t{}\n'.format(id2entity[h], id2relation[r], id2entity[t], s)) # Annotate hidden triplets print('annotation') cand = {} with gzip.open('/common/users/yz956/kg/code/KBRD/data/cpa/cpa/kg_test_candidates_Cell_Phones_and_Accessories.txt.gz', 'rt') as f: for line in f: cells = line.split() uid = int(cells[0]) item_ids = [int(i) for i in cells[1:]] cand[uid] = item_ids ann, train = [], [] d = {} with open('/common/users/yz956/kg/code/KBRD/data/cpa/cpa/sample_pre.txt') as ft: for line in ft: line = line.strip().split('\t') train.append(line[1:]) for u in range(61254): hiddens = [] for i in cand[u]: # for i in range(61254, 108858): hiddens.append((u, 0, i)) scores = kge_model.infer_step(kge_model, hiddens, args) score_np = np.array(scores) d = dict(zip(cand[u], scores)) # d = dict(zip(range(61254, 108858), scores)) d = sorted(d.items(), key=lambda x: x[1], reverse=True) # d_50 = d[:50] # for idx, t in enumerate(train[u]): # for (tt, prob) in d_50: # if int(t) == tt: # d_50.remove((tt, prob)) # d_50.append(d[50 + idx]) # assert len(d_50) == 50 # d = {} d_50 = d ann.append(d_50) with open(local_path + '/annotation_1000_htr.txt', 'w') as fo: for idx, d in enumerate(ann): for (t, score) in d: fo.write(str(idx) + '\t' + str(t) + '\t0\t' + str(score) + '\n') # with open(local_path + '/hidden_50_p.txt', 'w') as fo: # for idx, d in enumerate(ann): # for (t, score) in d: # fo.write(str(idx) + '\t' + str(t) + '\t0\n') scores = kge_model.infer_step(kge_model, hidden_triples, args) with open(local_path + '/annotation_htr.txt', 'w') as fo: for (h, r, t), s in zip(hidden_triples, scores): # fo.write('{}\t{}\t{}\t{}\n'.format(id2entity[h], id2relation[r], id2entity[t], s)) fo.write('{}\t{}\t{}\t{}\n'.format(str(h), str(t), str(r), s)) if args.evaluate_train: logging.info('Evaluating on Training Dataset...') metrics, preds = kge_model.test_step(kge_model, train_triples, all_true_triples, args) log_metrics('Test', step, metrics)
def main(args): #什么模式 if (not args.do_train) and (not args.do_valid) and (not args.do_test): raise ValueError('one of train/val/test mode must be choosed.') #是否要用config修改一些命令行参数 if args.init_checkpoint: # 如果init_checkpoint有值(是一个路径)就从config里面修改一些参数 override_config(args) #这个函数里面要用到init_checkpoint,因此进了这个if就说明它有值就不会报错 elif args.data_path is None: #override函数里,如果path是none就从config里读,如果就没进上一个if,就不能从config里读,默认是none,后来自己改成默认值,按说这里不指定的话就会raise这个错误 raise ValueError('one of init_checkpoint/data_path must be choosed.') #训练但是没给保存路径 if args.do_train and args.save_path is None: raise ValueError('Where do you want to save your trained model?') #有保存路径但是文件夹不存在,就创建一个 if args.save_path and not os.path.exists(args.save_path): os.makedirs(args.save_path) # Write logs to checkpoint and console set_logger(args) #里面是所有实体/关系的代码 1 xx 2 xx 3 xx #读出来的entity2id,relation2id是键为实体/关系代码,值为序号的字典 with open(os.path.join(args.data_path, 'entities.dict')) as fin: entity2id = dict() for line in fin: eid, entity = line.strip().split( '\t') #1 xx eid=str(1),entity=xx entity2id[entity] = int(eid) with open(os.path.join(args.data_path, 'relations.dict')) as fin: relation2id = dict() for line in fin: rid, relation = line.strip().split('\t') relation2id[relation] = int(rid) # Read regions for Countries S* datasets if args.countries: regions = list() #contries数据集里面会有一个文件是regions.list其他的数据集没有 with open(os.path.join(args.data_path, 'regions.list')) as fin: for line in fin: #只有5行 region = line.strip() regions.append(entity2id[region]) #这里的值应该是区域的序号 args.regions = regions nentity = len(entity2id) nrelation = len(relation2id) args.nentity = nentity args.nrelation = nrelation logging.info('Model: %s' % args.model) logging.info('Data Path: %s' % args.data_path) logging.info('#entity: %d' % nentity) logging.info('#relation: %d' % nrelation) #开始获取训练验证测试数据集,并打印size train_triples = read_triple(os.path.join(args.data_path, 'train.txt'), entity2id, relation2id) logging.info('#train: %d' % len(train_triples)) valid_triples = read_triple(os.path.join(args.data_path, 'valid.txt'), entity2id, relation2id) logging.info('#valid: %d' % len(valid_triples)) test_triples = read_triple(os.path.join(args.data_path, 'test.txt'), entity2id, relation2id) logging.info('#test: %d' % len(test_triples)) # All true triples all_true_triples = train_triples + valid_triples + test_triples #构造模型 kge_model = KGEModel( model_name=args.model, nentity=nentity, nrelation=nrelation, hidden_dim=args.hidden_dim, gamma=args.gamma, double_entity_embedding=args.double_entity_embedding, double_relation_embedding=args.double_relation_embedding) logging.info('Model Parameter Configuration:') for name, param in kge_model.named_parameters(): logging.info('Parameter %s: %s, require_grad = %s' % (name, str(param.size()), str(param.requires_grad))) if args.cuda: kge_model = kge_model.cuda() if args.do_train: # Set training dataloader iterator train_dataloader_head = DataLoader( TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, 'head-batch'), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn) train_dataloader_tail = DataLoader( TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, 'tail-batch'), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn) train_iterator = BidirectionalOneShotIterator(train_dataloader_head, train_dataloader_tail) # Set training configuration current_learning_rate = args.learning_rate optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, kge_model.parameters()), #fitter操作,只优化requires_grad为true的 lr=current_learning_rate) if args.warm_up_steps: warm_up_steps = args.warm_up_steps else: warm_up_steps = args.max_steps // 2 if args.init_checkpoint: # Restore model from checkpoint directory logging.info('Loading checkpoint %s...' % args.init_checkpoint) checkpoint = torch.load( os.path.join(args.init_checkpoint, 'checkpoint')) init_step = checkpoint['step'] kge_model.load_state_dict(checkpoint['model_state_dict']) if args.do_train: current_learning_rate = checkpoint['current_learning_rate'] warm_up_steps = checkpoint['warm_up_steps'] optimizer.load_state_dict(checkpoint['optimizer_state_dict']) else: logging.info('Ramdomly Initializing %s Model...' % args.model) init_step = 0 step = init_step logging.info('Start Training...') logging.info('init_step = %d' % init_step) logging.info('learning_rate = %d' % current_learning_rate) logging.info('batch_size = %d' % args.batch_size) logging.info('negative_adversarial_sampling = %d' % args.negative_adversarial_sampling) logging.info('hidden_dim = %d' % args.hidden_dim) logging.info('gamma = %f' % args.gamma) logging.info('negative_adversarial_sampling = %s' % str(args.negative_adversarial_sampling)) if args.negative_adversarial_sampling: logging.info('adversarial_temperature = %f' % args.adversarial_temperature) # Set valid dataloader as it would be evaluated during training if args.do_train: training_logs = [] # Training Loop for step in range(init_step, args.max_steps): #train_iterator = BidirectionalOneShotIterator(train_dataloader_head, train_dataloader_tail) log = kge_model.train_step(kge_model, optimizer, train_iterator, args) training_logs.append(log) #动态调整学习率 if step >= warm_up_steps: #大于warm_up_steps后学习率变为原来的1/10 current_learning_rate = current_learning_rate / 10 logging.info('Change learning_rate to %f at step %d' % (current_learning_rate, step)) optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, kge_model.parameters()), lr=current_learning_rate #更新优化器里的学习率 ) warm_up_steps = warm_up_steps * 3 #更新warm_up_steps #每隔save_checkpoint_steps保存一次模型 if step % args.save_checkpoint_steps == 0: save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, 'warm_up_steps': warm_up_steps } save_model(kge_model, optimizer, save_variable_list, args) if step % args.log_steps == 0: metrics = {} for metric in training_logs[0].keys(): metrics[metric] = sum( [log[metric] for log in training_logs]) / len(training_logs) log_metrics('Training average', step, metrics) training_logs = [] if args.do_valid and step % args.valid_steps == 0: logging.info('Evaluating on Valid Dataset...') metrics = kge_model.test_step(kge_model, valid_triples, all_true_triples, args) log_metrics('Valid', step, metrics) save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, 'warm_up_steps': warm_up_steps } save_model(kge_model, optimizer, save_variable_list, args) if args.do_valid: logging.info('Evaluating on Valid Dataset...') metrics = kge_model.test_step(kge_model, valid_triples, all_true_triples, args) log_metrics('Valid', step, metrics) if args.do_test: logging.info('Evaluating on Test Dataset...') metrics = kge_model.test_step(kge_model, test_triples, all_true_triples, args) log_metrics('Test', step, metrics) if args.evaluate_train: logging.info('Evaluating on Training Dataset...') metrics = kge_model.test_step(kge_model, train_triples, all_true_triples, args) log_metrics('Test', step, metrics)
def main(args): os.environ["CUDA_VISIBLE_DEVICES"] = args.CUDA_DEVISE args.data_path = os.path.join(args.datadir, args.dataset, 'onto_file') # if args.init_checkpoint: # override_config(args) if args.data_path is None: raise ValueError('data_path and dataset must be choosed.') args.save_path = os.path.join(args.data_path, 'save_onto_embeds') # if args.do_train and args.save_path is None: # raise ValueError('Where do you want to save your trained model?') if args.save_path and not os.path.exists(args.save_path): os.makedirs(args.save_path) with open(os.path.join(args.data_path, 'entities.dict')) as fin: entity2id = dict() for line in fin: eid, entity = line.strip().split('\t') entity2id[entity] = int(eid) with open(os.path.join(args.data_path, 'relations.dict')) as fin: relation2id = dict() for line in fin: rid, relation = line.strip().split('\t') relation2id[relation] = int(rid) nentity = len(entity2id) nrelation = len(relation2id) args.nentity = nentity args.nrelation = nrelation print('Model: %s' % args.model) # print('Data Path: %s' % args.data_path + "/" + args.dataset) print('#entity num: %d' % nentity) print('#relation num: %d' % nrelation) all_triples = read_triple(os.path.join(args.data_path, 'all_triples.txt'), entity2id, relation2id) print('#total triples num: %d' % len(all_triples)) # All true triples all_true_triples = all_triples kge_model = KGEModel( model_name=args.model, nentity=nentity, nrelation=nrelation, hidden_dim=args.hidden_dim, gamma=args.gamma, double_entity_embedding=args.double_entity_embedding, double_relation_embedding=args.double_relation_embedding) # logging.info('Model Parameter Configuration:') # for name, param in kge_model.named_parameters(): # logging.info('Parameter %s: %s, require_grad = %s' % (name, str(param.size()), str(param.requires_grad))) if args.cuda: kge_model = kge_model.cuda() if args.do_train: # Set training dataloader iterator train_dataloader_head = DataLoader( TrainDataset(all_triples, nentity, nrelation, args.negative_sample_size, 'head-batch'), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn) train_dataloader_tail = DataLoader( TrainDataset(all_triples, nentity, nrelation, args.negative_sample_size, 'tail-batch'), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn) train_iterator = BidirectionalOneShotIterator(train_dataloader_head, train_dataloader_tail) # Set training configuration current_learning_rate = args.learning_rate optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, kge_model.parameters()), lr=current_learning_rate) if args.warm_up_steps: warm_up_steps = args.warm_up_steps else: warm_up_steps = args.max_steps // 2 print('Ramdomly Initializing %s Model...' % args.model) # step = init_step print('------ Start Training...') print('batch_size = %d' % args.batch_size) print('negative sample size = %d' % args.negative_sample_size) print('hidden_dim = %d' % args.hidden_dim) print('gamma = %f' % args.gamma) print('negative_adversarial_sampling = %s' % str(args.negative_adversarial_sampling)) if args.negative_adversarial_sampling: print('adversarial_temperature = %f' % args.adversarial_temperature) print("learning rate = %f" % current_learning_rate) # Set valid dataloader as it would be evaluated during training if args.do_train: train_losses = [] # Training Loop for step in range(1, args.max_steps + 1): loss_values = kge_model.train_step(kge_model, optimizer, train_iterator, args) train_losses.append(loss_values) if step >= warm_up_steps: current_learning_rate = current_learning_rate / 10 print('Change learning_rate to %f at step %d' % (current_learning_rate, step)) optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, kge_model.parameters()), lr=current_learning_rate) warm_up_steps = warm_up_steps * 3 if step % args.print_steps == 0: pos_sample_loss = sum( [losses['pos_sample_loss'] for losses in train_losses]) / len(train_losses) neg_sample_loss = sum( [losses['neg_sample_loss'] for losses in train_losses]) / len(train_losses) loss1 = sum([losses['loss'] for losses in train_losses]) / len(train_losses) # log_metrics('Training average', step, metrics) print( 'Training Step: %d; average -> pos_sample_loss: %f; neg_sample_loss: %f; loss: %f' % (step, pos_sample_loss, neg_sample_loss, loss1)) train_losses = [] if step % args.save_steps == 0: save_embeddings(kge_model, step, args) if args.evaluate_train and step % args.valid_steps == 0: print('------ Evaluating on Training Dataset...') metrics = kge_model.test_step(kge_model, all_triples, all_true_triples, args) log_metrics('Test', step, metrics)
def main(args): if (not args.do_train) and (not args.do_valid) and (not args.do_test): raise ValueError('one of train/val/test mode must be choosed.') if args.init_checkpoint: override_config(args) elif args.data_path is None: raise ValueError('one of init_checkpoint/data_path must be choosed.') ''' #if need to store the outputs elif args.storefile is None: raise ValueError('where do you want to store you result?') ''' if args.do_train and args.save_path is None: raise ValueError('Where do you want to save your trained model?') if args.save_path and not os.path.exists(args.save_path): os.makedirs(args.save_path) # Write logs to checkpoint and console set_logger(args) with open(os.path.join(args.data_path, 'entities.txt')) as fin: all_entitys = list() for line in fin: entity, eid = line.strip().split('\t') all_entitys.append(eid) #define the entity file and inputdataWWC with open(os.path.join(args.data_path, 'entities.txt')) as fin: entity2id = dict() for line in fin: entity, eid = line.strip().split('\t') entity2id[entity] = int(eid) nentity = len(entity2id) args.nentity = nentity logging.info('Model: %s' % args.model) logging.info('Data Path: %s' % args.data_path) logging.info('#entity: %d' % nentity) #laod all dataset train_triples = read_triple(os.path.join(args.data_path, 'train.txt'), entity2id) logging.info('#train: %d' % len(train_triples)) valid_triples = read_triple(os.path.join(args.data_path, 'valid.txt'), entity2id) logging.info('#valid: %d' % len(valid_triples)) test_triples = read_triple(os.path.join(args.data_path, 'test.txt'), entity2id) logging.info('#test: %d' % len(test_triples)) addtrain_triples = read_triple(os.path.join(args.data_path, 'test.txt'), entity2id) logging.info('#add_train:%d' % len(addtrain_triples)) #All true triples all_true_triples = train_triples + valid_triples + test_triples + addtrain_triples train_triples = train_triples #+ addtrain_triples if need experiment,delete addtrain ''' #load output mode dataset train_triples = read_triple(os.path.join(args.data_path, 'train_origin.txt'), entity2id, relation2id) logging.info('#train: %d' % len(train_triples)) valid_triples = read_triple(os.path.join(args.data_path, 'valid_1.txt'), entity2id, relation2id) logging.info('#valid: %d' % len(valid_triples)) test_triples = read_triple(os.path.join(args.data_path, 'test_1.txt'), entity2id, relation2id) logging.info('#test: %d' % len(test_triples)) #All true triples all_true_triples = train_triples + valid_triples + test_triples train_triples = train_triples + valid_triples ''' bdg_model = Bridge_rules( nentity=nentity, hidden_dim=args.hidden_dim, gamma=args.gamma, ) logging.info('Model Parameter Configuration:') for name, param in bdg_model.named_parameters(): logging.info('Parameter %s: %s, require_grad = %s' % (name, str(param.size()), str(param.requires_grad))) if args.cuda: bdg_model = bdg_model.cuda() if args.do_train: # Set training dataloader iterator train_dataloader_head = DataLoader( TrainDataset(train_triples, nentity, args.negative_sample_size, 'head-batch'), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn) train_dataloader_tail = DataLoader( TrainDataset(train_triples, nentity, args.negative_sample_size, 'tail-batch'), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn) train_iterator = BidirectionalOneShotIterator(train_dataloader_head, train_dataloader_tail) # Set training configuration current_learning_rate = args.learning_rate optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, bdg_model.parameters()), lr=current_learning_rate) if args.warm_up_steps: warm_up_steps = args.warm_up_steps else: warm_up_steps = args.max_steps // 2 if args.init_checkpoint: # Restore model from checkpoint directory logging.info('Loading checkpoint %s...' % args.init_checkpoint) checkpoint = torch.load( os.path.join(args.init_checkpoint, 'checkpoint')) init_step = checkpoint['step'] bdg_model.load_state_dict(checkpoint['model_state_dict']) if args.do_train: current_learning_rate = checkpoint['current_learning_rate'] warm_up_steps = checkpoint['warm_up_steps'] optimizer.load_state_dict(checkpoint['optimizer_state_dict']) else: logging.info('Ramdomly Initializing %s Model...' % args.model) init_step = 0 step = init_step logging.info('Start Training...') logging.info('init_step = %d' % init_step) logging.info('batch_size = %d' % args.batch_size) logging.info('negative_adversarial_sampling = %d' % args.negative_adversarial_sampling) logging.info('hidden_dim = %d' % args.hidden_dim) logging.info('gamma = %f' % args.gamma) logging.info('negative_adversarial_sampling = %s' % str(args.negative_adversarial_sampling)) if args.negative_adversarial_sampling: logging.info('adversarial_temperature = %f' % args.adversarial_temperature) # Set valid dataloader as it would be evaluated during training if args.do_train: logging.info('learning_rate = %d' % current_learning_rate) training_logs = [] #Training Loop for step in range(init_step, args.max_steps): log = bdg_model.train_step(bdg_model, optimizer, train_iterator, args) training_logs.append(log) if step >= warm_up_steps: current_learning_rate = current_learning_rate / 10 logging.info('Change learning_rate to %f at step %d' % (current_learning_rate, step)) optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, bdg_model.parameters()), lr=current_learning_rate) warm_up_steps = warm_up_steps * 3 if step % args.save_checkpoint_steps == 0: save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, 'warm_up_steps': warm_up_steps } save_model(bdg_model, optimizer, save_variable_list, args) if step % args.log_steps == 0: metrics = {} for metric in training_logs[0].keys(): metrics[metric] = sum( [log[metric] for log in training_logs]) / len(training_logs) log_metrics('Training average', step, metrics) training_logs = [] if args.do_valid and step % args.valid_steps == 0: logging.info('Evaluating on Valid Dataset...') metrics = bdg_model.test_step(bdg_model, valid_triples, all_true_triples, args) log_metrics('Valid', step, metrics) save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, 'warm_up_steps': warm_up_steps } save_model(bdg_model, optimizer, save_variable_list, args) if args.do_valid: logging.info('Evaluating on Valid Dataset...') metrics = bdg_model.test_step(bdg_model, valid_triples, all_true_triples, args) log_metrics('Valid', step, metrics) if args.do_test: logging.info('Evaluating on Test Dataset...') metrics = bdg_model.test_step(bdg_model, test_triples, all_true_triples, args) log_metrics('Test', step, metrics) if args.evaluate_train: logging.info('Evaluating on Training Dataset...') metrics = bdg_model.test_step(bdg_model, train_triples, all_true_triples, args) log_metrics('Test', step, metrics) if args.output: logging.info('Output the entities rank...') storefile = 'score.txt' bdg_model.store_score(bdg_model, test_triples, all_true_triples, storefile, args)
def train(self,train_triples,time,entity2id,relation2id,valid_triples): self.nentity = len(entity2id) self.nrelation = len(relation2id) current_learning_rate = 0.0001 if time == -1: # self.kge_model = KGEModel( # ent_tot=self.nentity, # rel_tot=self.nrelation, # dim_e=50, # dim_r=50 # ) # self.kge_model = self.kge_model.cuda() # self.optimizer = torch.optim.Adam( # filter(lambda p: p.requires_grad, self.kge_model.parameters()), # lr=current_learning_rate # ) # init_step = 0 #else: # temp = time - 1 self.kge_model, self.optimizer, step_loaded = self.load_model(time) init_step = step_loaded else: init_step = 0 # Set training dataloader iterator train_dataloader_head = DataLoader( TrainDataset(train_triples, self.nentity, self.nrelation, round(len(train_triples)*0.1), 'head-batch'), batch_size=50, shuffle=False, num_workers=2, collate_fn=TrainDataset.collate_fn ) train_dataloader_tail = DataLoader( TrainDataset(train_triples, self.nentity, self.nrelation, round(len(train_triples)*0.1), 'tail-batch'), batch_size=50, shuffle=False, num_workers=2, collate_fn=TrainDataset.collate_fn ) warm_up_steps = 5000 // 2 train_iterator = BidirectionalOneShotIterator(train_dataloader_head, train_dataloader_tail) if self.isCUDA == 1: self.kge_model = self.kge_model.cuda() #start training print("start training:%d"%time) # Training Loop starttime = Time.time() if time==-1: steps = 3000*4 printnum = 100 else: steps = 150 printnum = 50 for step in range(init_step, steps): loss = self.kge_model.train_step(self.kge_model, self.optimizer, train_iterator,self.isCUDA) ''' if step >= warm_up_steps: current_learning_rate = current_learning_rate / 10 logging.info('Change learning_rate to %f at step %d' % (current_learning_rate, step)) self.optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, self.kge_model.parameters()), lr=current_learning_rate ) warm_up_steps = warm_up_steps * 3 ''' if step%printnum==0: endtime = Time.time() print("step:%d, cost time: %s, loss is %.4f" % (step,round((endtime - starttime), 3),loss)) '''self.save_model(self.kge_model, self.optimizer, time) result_head, result_tail = self.evaluate(valid_triples, train_triples + valid_triples, relation2id, entity2id, time) print(result_head) print(result_tail) self.kge_model, self.optimizer = self.load_model(time)''' self.save_model(self.kge_model, self.optimizer, time, step)
def main(args): if not torch.cuda.is_available(): args.cuda = False if args.ruge: args.loss = 'ruge' if (not args.do_train) and (not args.do_valid) and (not args.do_test) and ( not args.do_experiment) and (not args.do_grid): raise ValueError('one of train/val/test mode must be choosed.') if args.init_checkpoint: override_config(args) elif args.data_path is None: raise ValueError('one of init_checkpoint/data_path must be choosed.') if args.do_train and args.save_path is None: raise ValueError('Where do you want to save your trained model?') if args.save_path and not os.path.exists(args.save_path): os.makedirs(args.save_path) # Write logs to checkpoint and console set_logger(args) if args.regularization != 0: print('L3 regularization with coeff - ', args.regularization) if args.l2_r != 0: print('L2 regularization with coeff - ', args.l2_r) if args.project != 0: print('projecting before training') #logging.info('Inverse loss = premise - concl (reverse)') if OPT_STOPPING: logging.info('Opt stopping is ON') print('Opt stopping is on') with open(os.path.join(args.data_path, 'entities.dict')) as fin: entity2id = dict() for line in fin: eid, entity = line.strip().split('\t') entity2id[entity] = int(eid) with open(os.path.join(args.data_path, 'relations.dict')) as fin: relation2id = dict() for line in fin: rid, relation = line.strip().split('\t') relation2id[relation] = int(rid) # Read regions for Countries S* datasets if args.countries: regions = list() with open(os.path.join(args.data_path, 'regions.list')) as fin: for line in fin: region = line.strip() regions.append(entity2id[region]) args.regions = regions nentity = len(entity2id) nrelation = len(relation2id) args.nentity = nentity args.nrelation = nrelation if args.inject: logging.info('With rule injection') else: logging.info('NO INJECTION') logging.info('Model: %s' % args.model) logging.info('Data Path: %s' % args.data_path) logging.info('#entity: %d' % nentity) logging.info('#relation: %d' % nrelation) train_triples = read_triple(os.path.join(args.data_path, 'train.txt'), entity2id, relation2id) logging.info('#train: %d' % len(train_triples)) valid_triples = read_triple(os.path.join(args.data_path, 'valid.txt'), entity2id, relation2id) logging.info('#valid: %d' % len(valid_triples)) test_triples = read_triple( os.path.join(args.data_path, 'test.txt'), entity2id, relation2id) # For testing on Symmetric in WordNet: Symmetric_testWN18 logging.info('#test: %d' % len(test_triples)) #All true triples all_true_triples = train_triples + valid_triples + test_triples train_args = {} # set up rule iterators device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') n_batches = len(train_triples) // args.batch_size if n_batches < len(train_triples) / args.batch_size: n_batches += 1 rule_iterators = {} rules_info = '' if args.inv: n_inverse, inverse_batchsize, rule_iterators[ 'inverse'] = setup_rule_loader(n_batches, args.batch_size, args.data_path, 'groundings_inverse.txt', device, RULE_BATCH_SIZE_INV) rules_info += 'Inverse: batch size %d out of %d rules' % ( inverse_batchsize, n_inverse) + '\n' if args.eq: n_eq, eq_batchsize, rule_iterators['equality'] = setup_rule_loader( n_batches, args.batch_size, args.data_path, 'groundings_equality.txt', device, RULE_BATCH_SIZE_EQ) rules_info += 'Equality: batch size %d out of %d rules' % ( eq_batchsize, n_eq) + '\n' if args.impl: n_impl, impl_batchsize, rule_iterators[ 'implication'] = setup_rule_loader(n_batches, args.batch_size, args.data_path, 'groundings_implication.txt', device, RULE_BATCH_SIZE_IMPL) rules_info += 'implication: batch size %d out of %d rules\n' % ( impl_batchsize, n_impl) if args.sym: n_symmetry, sym_batchsize, rule_iterators[ 'symmetry'] = setup_rule_loader(n_batches, args.batch_size, args.data_path, 'groundings_symmetric.txt', device, RULE_BATCH_SIZE_SYM) rules_info += 'symmetry: batch size %d out of %d rules\n' % ( sym_batchsize, n_symmetry) if args.ruge or args.ruge_inject: n_rules, rule_iterators['ruge'] = construct_ruge_loader( n_batches, args) rules_info += 'RUGE: Total %d rules\n' % n_rules if rules_info: logging.info(rules_info) # ----------- adversarial ------------------ if args.adversarial: clauses_filename = os.path.join(args.data_path, 'clauses_0.9.pl') adv_clauses, clentity2id = dt.read_clauses(clauses_filename, relation2id) n_clause_entities = len(clentity2id) mult = 2 if args.model in ['TransE', 'pRotatE']: mult = 1 if 'QuatE' in args.model: mult = 4 adv_model = ADVModel(clauses=adv_clauses, n_entities=len(clentity2id), dim=mult * args.hidden_dim, use_cuda=args.cuda) if args.cuda: adv_model = adv_model.cuda() else: adv_model = None if args.do_grid: if rules_info: print(rules_info) run_grid(nentity, nrelation, train_triples, valid_triples, test_triples, all_true_triples, args, rule_iterators, adv_model) exit() ntriples = len(train_triples) kge_model = KGEModel( model_name=args.model, nentity=nentity, nrelation=nrelation, ntriples=ntriples, hidden_dim=args.hidden_dim, gamma=args.gamma, ) kge_model.set_loss(args.loss) logging.info('Model Parameter Configuration:') for name, param in kge_model.named_parameters(): logging.info('Parameter %s: %s, require_grad = %s' % (name, str(param.size()), str(param.requires_grad))) logging.info('Loss function %s' % args.loss) if args.cuda and args.parallel: gpus = [0, 1] os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(str(x) for x in gpus) kge_model.cuda() kge_model = torch.nn.DataParallel(kge_model, device_ids=[0, 1]) elif args.cuda: kge_model = kge_model.cuda() if args.do_train or args.do_experiment: # Set training dataloader iterator train_dataloader_head = DataLoader( TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, 'head-batch'), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn) train_dataloader_tail = DataLoader( TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, 'tail-batch'), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn) train_iterator = BidirectionalOneShotIterator(train_dataloader_head, train_dataloader_tail) # Set training configuration current_learning_rate = args.learning_rate optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, kge_model.parameters()), lr=current_learning_rate) if args.init_checkpoint: # Restore model from checkpoint directory logging.info('Loading checkpoint %s...' % args.init_checkpoint) checkpoint = torch.load( os.path.join(args.init_checkpoint, 'checkpoint')) init_step = checkpoint['step'] kge_model.load_state_dict(checkpoint['model_state_dict']) if args.do_train: current_learning_rate = checkpoint['current_learning_rate'] warm_up_steps = checkpoint['warm_up_steps'] optimizer.load_state_dict(checkpoint['optimizer_state_dict']) else: logging.info('Ramdomly Initializing %s Model...' % args.model) init_step = 0 step = init_step logging.info('Start Training...') logging.info('init_step = %d' % init_step) logging.info('learning_rate = %d' % current_learning_rate) logging.info('batch_size = %d' % args.batch_size) logging.info('negative_adversarial_sampling = %d' % args.negative_adversarial_sampling) logging.info('hidden_dim = %d' % args.hidden_dim) logging.info('gamma = %f' % args.gamma) logging.info('negative_adversarial_sampling = %s' % str(args.negative_adversarial_sampling)) if args.negative_adversarial_sampling: logging.info('adversarial_temperature = %f' % args.adversarial_temperature) # Set valid dataloader as it would be evaluated during training if args.do_train: train_model(init_step, valid_triples, all_true_triples, kge_model, train_iterator, len(train_triples), args) if args.evaluate_train: logging.info('Evaluating on Training Dataset...') model_module = kge_model.module if args.parallel else kge_model metrics = model_module.test_step(kge_model, train_triples, all_true_triples, args) #metrics1 = model_module.getScore(kge_model, train_triples, all_true_triples, args) log_metrics('Test', step, metrics) # experiment on the updated function if args.do_experiment: logging.info('\n\nSTARTING EXPERIMENT\n') train_model(init_step, valid_triples, all_true_triples, kge_model, train_iterator, rule_iterators, args) if args.do_valid: logging.info('Evaluating on Valid Dataset...') model_module = kge_model.module if args.parallel else kge_model metrics = model_module.test_step(kge_model, valid_triples, all_true_triples, args) #metrics1 = model_module.getScore(kge_model, train_triples, all_true_triples, args) log_metrics('Valid', step, metrics) if args.do_test: logging.info('Evaluating on Test Dataset...') model_module = kge_model.module if args.parallel else kge_model metrics = model_module.test_step(kge_model, test_triples, all_true_triples, args) log_metrics('Test', step, metrics) if args.evaluate_train: logging.info('Evaluating on Training Dataset...') model_module = kge_model.module if args.parallel else kge_model metrics = model_module.test_step(kge_model, train_triples, all_true_triples, args) log_metrics('Test', step, metrics)
def main(args): if (not args.do_train) and (not args.do_valid) and (not args.do_test): raise ValueError('one of train/val/test mode must be choosed.') if not args.do_train and args.init_checkpoint: override_config(args) elif args.data_path is None: raise ValueError('one of init_checkpoint/data_path must be choosed.') if args.do_train and args.save_path is None: raise ValueError('Where do you want to save your trained model?') if args.save_path and not os.path.exists(args.save_path): os.makedirs(args.save_path) # Write logs to checkpoint and console set_logger(args) with open(os.path.join(args.data_path, 'entities.dict')) as fin: entity2id = dict() for line in fin: eid, entity = line.strip().split('\t') entity2id[entity] = int(eid) with open(os.path.join(args.data_path, 'relations.dict')) as fin: relation2id = dict() for line in fin: rid, relation = line.strip().split('\t') relation2id[relation] = int(rid) nentity = len(entity2id) nrelation = len(relation2id) args.nentity = nentity args.nrelation = nrelation logging.info('Model: %s' % args.model) logging.info('Data Path: %s' % args.data_path) logging.info('#entity: %d' % nentity) logging.info('#relation: %d' % nrelation) train_triples = read_triple(os.path.join(args.data_path, "train.txt"), entity2id, relation2id) if args.self_test: train_triples = train_triples[len(train_triples) // 5:] if args.fake: fake_triples = pickle.load( open(os.path.join(args.data_path, "fake%s.pkl" % args.fake), "rb")) fake = torch.LongTensor(fake_triples) train_triples += fake_triples else: fake_triples = [(0, 0, 0)] fake = torch.LongTensor(fake_triples) if args.cuda: fake = fake.cuda() logging.info('#train: %d' % len(train_triples)) valid_triples = read_triple(os.path.join(args.data_path, 'valid.txt'), entity2id, relation2id) logging.info('#valid: %d' % len(valid_triples)) test_triples = read_triple(os.path.join(args.data_path, 'test.txt'), entity2id, relation2id) logging.info('#test: %d' % len(test_triples)) all_true_triples = train_triples + valid_triples + test_triples kge_model = KGEModel( model_name=args.model, nentity=nentity, nrelation=nrelation, hidden_dim=args.hidden_dim, gamma=args.gamma, double_entity_embedding=args.double_entity_embedding, double_relation_embedding=args.double_relation_embedding) trainer = None if args.method == "CLF": trainer = ClassifierTrainer(train_triples, fake_triples, args, kge_model, args.hard) elif args.method == "LT": trainer = LTTrainer(train_triples, fake_triples, args, kge_model) elif args.method == "NoiGAN": trainer = NoiGANTrainer(train_triples, fake_triples, args, kge_model, args.hard) logging.info('Model Parameter Configuration:') for name, param in kge_model.named_parameters(): logging.info('Parameter %s: %s, require_grad = %s' % (name, str(param.size()), str(param.requires_grad))) if args.cuda: kge_model = kge_model.cuda() if args.do_train: # Set training dataloader iterator train_dataloader_head = DataLoader( TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, 'head-batch'), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn) train_dataloader_tail = DataLoader( TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, 'tail-batch'), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn) train_iterator = BidirectionalOneShotIterator(train_dataloader_head, train_dataloader_tail) # Set training configuration current_learning_rate = args.learning_rate optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, kge_model.parameters()), lr=current_learning_rate) if args.warm_up_steps: warm_up_steps = args.warm_up_steps else: warm_up_steps = args.max_steps if args.init_checkpoint: # Restore model from checkpoint directory logging.info('Loading checkpoint %s...' % args.init_checkpoint) checkpoint = torch.load( os.path.join(args.init_checkpoint, 'checkpoint')) init_step = 0 #checkpoint['step'] kge_model.load_state_dict(checkpoint['model_state_dict']) if args.do_train: # current_learning_rate = checkpoint['current_learning_rate'] # warm_up_steps = checkpoint['warm_up_steps'] optimizer.load_state_dict(checkpoint['optimizer_state_dict']) else: logging.info('Ramdomly Initializing %s Model...' % args.model) init_step = 0 step = init_step logging.info('Start Training...') logging.info('init_step = %d' % init_step) logging.info('batch_size = %d' % args.batch_size) logging.info('negative_adversarial_sampling = %d' % args.negative_adversarial_sampling) logging.info('hidden_dim = %d' % args.hidden_dim) logging.info('gamma = %f' % args.gamma) logging.info('negative_adversarial_sampling = %s' % str(args.negative_adversarial_sampling)) if args.negative_adversarial_sampling: logging.info('adversarial_temperature = %f' % args.adversarial_temperature) # Set valid dataloader as it would be evaluated during training if args.do_train: logging.info('learning_rate = %f' % current_learning_rate) training_logs = [] #Training Loop triple2confidence_weights = None for step in range(init_step, args.max_steps): if args.method == "CLF" and step % args.classify_steps == 0: logging.info('Train Classifier') metrics = trainer.train_classifier(trainer) log_metrics('Classifier', step, metrics) metrics = trainer.test_ave_score(trainer) log_metrics('Classifier', step, metrics) trainer.cal_confidence_weight() elif args.method == "NoiGAN" and step % args.classify_steps == 0: logging.info('Train NoiGAN') trainer.train_NoiGAN(trainer) metrics = trainer.test_ave_score(trainer) log_metrics('Classifier', step, metrics) trainer.cal_confidence_weight() log = kge_model.train_step(kge_model, optimizer, train_iterator, args, trainer=trainer) training_logs.append(log) if step >= warm_up_steps: current_learning_rate = current_learning_rate / 10 logging.info('Change learning_rate to %f at step %d' % (current_learning_rate, step)) optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, kge_model.parameters()), lr=current_learning_rate) warm_up_steps = warm_up_steps * 3 if step % args.save_checkpoint_steps == 0: save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, 'warm_up_steps': warm_up_steps } save_model(kge_model, optimizer, save_variable_list, args, trainer) if step % args.log_steps == 0: metrics = {} for metric in training_logs[0].keys(): metrics[metric] = sum( [log[metric] for log in training_logs]) / len(training_logs) log_metrics('Training average', step, metrics) training_logs = [] if args.do_valid and step % args.valid_steps == 0: logging.info('Evaluating on Valid Dataset...') metrics = kge_model.test_step(kge_model, valid_triples, all_true_triples, args) log_metrics('Valid', step, metrics) save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, 'warm_up_steps': warm_up_steps } save_model(kge_model, optimizer, save_variable_list, args, trainer) if trainer is not None: logging.info("Evaluating Classifier on Train Dataset") metrics = trainer.test_ave_score(trainer) log_metrics('Train', step, metrics) if args.do_valid: logging.info('Evaluating on Valid Dataset...') metrics = kge_model.test_step(kge_model, valid_triples, all_true_triples, args) log_metrics('Valid', step, metrics) if args.do_test: logging.info('Evaluating on Test Dataset...') metrics = kge_model.test_step(kge_model, test_triples, all_true_triples, args) log_metrics('Test', step, metrics) # logging.info("\t".join([metric for metric in metrics.values()])) if args.evaluate_train: logging.info('Evaluating on Training Dataset...') metrics = kge_model.test_step(kge_model, train_triples, all_true_triples, args) log_metrics('Test', step, metrics)
def main(args): if ( (not args.do_train) and (not args.do_valid) and (not args.do_test) and (not args.evaluate_train) ): raise ValueError("one of train/val/test mode must be choosed.") if args.init_checkpoint: override_config(args) args.save_path = ( "log/%s/%s/%s-%s/%s" % (args.dataset, args.model, args.hidden_dim, args.gamma, time.time()) if args.save_path == None else args.save_path ) writer = SummaryWriter(args.save_path) # Write logs to checkpoint and console set_logger(args) dataset = LinkPropPredDataset(name=args.dataset) split_dict = dataset.get_edge_split() nentity = dataset.graph["num_nodes"] nrelation = int(max(dataset.graph["edge_reltype"])[0]) + 1 evaluator = Evaluator(name=args.dataset) args.nentity = nentity args.nrelation = nrelation logging.info("Model: %s" % args.model) logging.info("Dataset: %s" % args.dataset) logging.info("#entity: %d" % nentity) logging.info("#relation: %d" % nrelation) train_triples = split_dict["train"] logging.info("#train: %d" % len(train_triples["head"])) valid_triples = split_dict["valid"] logging.info("#valid: %d" % len(valid_triples["head"])) test_triples = split_dict["test"] logging.info("#test: %d" % len(test_triples["head"])) train_count, train_true_head, train_true_tail = ( defaultdict(lambda: 4), defaultdict(list), defaultdict(list), ) for i in tqdm(range(len(train_triples["head"]))): head, relation, tail = ( train_triples["head"][i], train_triples["relation"][i], train_triples["tail"][i], ) train_count[(head, relation)] += 1 train_count[(tail, -relation - 1)] += 1 train_true_head[(relation, tail)].append(head) train_true_tail[(head, relation)].append(tail) kge_model = KGEModel( model_name=args.model, nentity=nentity, nrelation=nrelation, hidden_dim=args.hidden_dim, gamma=args.gamma, double_entity_embedding=args.double_entity_embedding, double_relation_embedding=args.double_relation_embedding, evaluator=evaluator, ) logging.info("Model Parameter Configuration:") for name, param in kge_model.named_parameters(): logging.info( "Parameter %s: %s, require_grad = %s" % (name, str(param.size()), str(param.requires_grad)) ) if args.cuda: kge_model = kge_model.cuda() if args.do_train: # Set training dataloader iterator train_dataloader_head = DataLoader( TrainDataset( train_triples, nentity, nrelation, args.negative_sample_size, "head-batch", train_count, train_true_head, train_true_tail, ), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn, ) train_dataloader_tail = DataLoader( TrainDataset( train_triples, nentity, nrelation, args.negative_sample_size, "tail-batch", train_count, train_true_head, train_true_tail, ), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn, ) train_iterator = BidirectionalOneShotIterator( train_dataloader_head, train_dataloader_tail ) # Set training configuration current_learning_rate = args.learning_rate optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, kge_model.parameters()), lr=current_learning_rate, ) if args.warm_up_steps: warm_up_steps = args.warm_up_steps else: warm_up_steps = args.max_steps // 2 if args.init_checkpoint: # Restore model from checkpoint directory logging.info("Loading checkpoint %s..." % args.init_checkpoint) checkpoint = torch.load(os.path.join(args.init_checkpoint, "checkpoint")) init_step = checkpoint["step"] kge_model.load_state_dict(checkpoint["model_state_dict"]) if args.do_train: current_learning_rate = checkpoint["current_learning_rate"] warm_up_steps = checkpoint["warm_up_steps"] optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) else: logging.info("Ramdomly Initializing %s Model..." % args.model) init_step = 0 step = init_step logging.info("Start Training...") logging.info("init_step = %d" % init_step) logging.info("batch_size = %d" % args.batch_size) logging.info( "negative_adversarial_sampling = %d" % args.negative_adversarial_sampling ) logging.info("hidden_dim = %d" % args.hidden_dim) logging.info("gamma = %f" % args.gamma) logging.info( "negative_adversarial_sampling = %s" % str(args.negative_adversarial_sampling) ) if args.negative_adversarial_sampling: logging.info("adversarial_temperature = %f" % args.adversarial_temperature) # Set valid dataloader as it would be evaluated during training if args.do_train: logging.info("learning_rate = %d" % current_learning_rate) training_logs = [] # Training Loop for step in range(init_step, args.max_steps): log = kge_model.train_step(kge_model, optimizer, train_iterator, args) training_logs.append(log) if step >= warm_up_steps: current_learning_rate = current_learning_rate / 10 logging.info( "Change learning_rate to %f at step %d" % (current_learning_rate, step) ) optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, kge_model.parameters()), lr=current_learning_rate, ) warm_up_steps = warm_up_steps * 3 if ( step % args.save_checkpoint_steps == 0 and step > 0 ): # ~ 41 seconds/saving save_variable_list = { "step": step, "current_learning_rate": current_learning_rate, "warm_up_steps": warm_up_steps, } save_model(kge_model, optimizer, save_variable_list, args) if step % args.log_steps == 0: metrics = {} for metric in training_logs[0].keys(): metrics[metric] = sum([log[metric] for log in training_logs]) / len( training_logs ) log_metrics("Train", step, metrics, writer) training_logs = [] if args.do_valid and step % args.valid_steps == 0 and step > 0: logging.info("Evaluating on Valid Dataset...") metrics = kge_model.test_step(kge_model, valid_triples, args) log_metrics("Valid", step, metrics, writer) save_variable_list = { "step": step, "current_learning_rate": current_learning_rate, "warm_up_steps": warm_up_steps, } save_model(kge_model, optimizer, save_variable_list, args) if args.do_valid: logging.info("Evaluating on Valid Dataset...") metrics = kge_model.test_step(kge_model, valid_triples, args) log_metrics("Valid", step, metrics, writer) if args.do_test: logging.info("Evaluating on Test Dataset...") metrics = kge_model.test_step(kge_model, test_triples, args) log_metrics("Test", step, metrics, writer) if args.evaluate_train: logging.info("Evaluating on Training Dataset...") small_train_triples = {} indices = np.random.choice( len(train_triples["head"]), args.ntriples_eval_train, replace=False ) for i in train_triples: small_train_triples[i] = train_triples[i][indices] metrics = kge_model.test_step( kge_model, small_train_triples, args, random_sampling=True ) log_metrics("Train", step, metrics, writer)
train_triples = read_triple("data/fr_en/att_triple_all") train_dataloader_head = data.DataLoader( TrainDataset(train_triples, entity_count, attr_count, value_count, 256, 'head-batch'), batch_size=1024, shuffle=False, num_workers=4, collate_fn=TrainDataset.collate_fn ) train_dataloader_tail = data.DataLoader( TrainDataset(train_triples, entity_count, attr_count, value_count, 256, 'tail-batch'), batch_size=1024, shuffle=False, num_workers=4, collate_fn=TrainDataset.collate_fn ) train_iterator = BidirectionalOneShotIterator(train_dataloader_head, train_dataloader_tail) model = TransE(entity_count, attr_count, value_count, device).to(device) optimizer = optim.SGD(model.parameters(), lr=learning_rate) start_epoch_id = 1 step = 0 best_score = 0.0 epochs = 8000 # if checkpoint_path: # start_epoch_id, step, best_score = load_checkpoint(checkpoint_path, model, optimizer) print(model) t = Tester()
def main(args): # fl == 0 : entity graph if args.fl == 0: if (not args.do_train_en) and (not args.do_valid_en) and ( not args.do_test_en): raise ValueError('One of train/valid/test mode must be choosed.') if args.init_checkpoint_en: override_config(args) elif args.data_path is None: raise ValueError( 'One of init_checkpoint/data_path must be choosed.') if args.do_train_en and args.save_path_en is None: raise ValueError('Where do you want to save your trained model?') if args.save_path_en and not os.path.exists(args.save_path_en): os.makedirs(args.save_path_en) set_logger(args) # 获取所有实体 final_entity.txt entity2id = dict() with open(os.path.join(args.data_path, 'final_entity_order.txt'), encoding='utf8') as fin: for line in fin.readlines(): eid, entity = line.strip().split('\t') entity2id[entity] = int(eid) nentity = len(entity2id) # 获取所有entity_relation 从ffinal_en_relation_order.txt entity_relation2id = dict() with open(os.path.join(args.data_path, 'ffinal_en_relation_order.txt')) as fin: for line in fin: en_reid, entity_relation = line.strip().split('\t') entity_relation2id[entity_relation] = int(en_reid) nentity_relation = len(entity_relation2id) args.nentity = nentity args.nentity_re = nentity_relation logging.info('entity_Model: %s' % args.entity_model) logging.info('Data Path: %s' % args.data_path) logging.info('number of entity: %d' % nentity) logging.info('number of entity_re: %d' % nentity_relation) #加载训练三元组集合, train_entity_triples = read_triple( os.path.join(args.data_path, 'train_entity_Graph.txt'), entity2id, entity_relation2id) logging.info('#train en_en triples: %s' % len(train_entity_triples)) #加载valid三元组集合, val_entity_triples = read_triple( os.path.join(args.data_path, 'val_entity_Graph.txt'), entity2id, entity_relation2id) logging.info('#val en_en triples: %s' % len(val_entity_triples)) #加载test三元组集合, test_entity_triples = read_triple( os.path.join(args.data_path, 'test_entity_Graph.txt'), entity2id, entity_relation2id) logging.info('#test en_en triples: %s' % len(test_entity_triples)) # all_triples all_true_triples_entity = train_entity_triples + val_entity_triples + test_entity_triples # entity model entity_kge_model = DKGE_Model( model_name=args.entity_model, nnode=args.nentity, nnode_re=args.nentity_re, hidden_dim=args.hidden_dim_en, gamma=args.gamma_en, gamma_intra=args.gamma_intra, double_node_embedding=args.double_node_embedding_en, double_node_re_embedding=args.double_node_re_embedding_en) logging.info('Entity Model Parameter Configuration:') for name, param in entity_kge_model.named_parameters(): logging.info('Parameter %s: %s, require_grad = %s' % (name, str(param.size()), str(param.requires_grad))) if args.cuda: entity_kge_model = entity_kge_model.cuda() if args.do_train_en: train_dataloader_head_en = DataLoader( TrainDataset_en(train_entity_triples, nentity, nentity_relation, args.negative_sample_size_en, 'head-batch'), batch_size=args.batch_size_en, shuffle=True, num_workers=int(max(1, args.cpu_num // 2)), collate_fn=TrainDataset_en.collate_fn) train_dataloader_tail_en = DataLoader( TrainDataset_en(train_entity_triples, nentity, nentity_relation, args.negative_sample_size_en, 'tail-batch'), batch_size=args.batch_size_en, shuffle=True, num_workers=int(max(1, args.cpu_num // 2)), collate_fn=TrainDataset_en.collate_fn) train_iterator_en = BidirectionalOneShotIterator( train_dataloader_head_en, train_dataloader_tail_en) # Set training configuration current_learning_rate_en = args.learning_rate_en optimizer_en = torch.optim.Adam(filter( lambda p: p.requires_grad, entity_kge_model.parameters()), lr=current_learning_rate_en, amsgrad=True) if args.warm_up_steps_en: warm_up_steps_en = args.warm_up_steps_en else: warm_up_steps_en = args.max_steps_en // 2 if args.init_checkpoint_en: # Restore model from checkpoint directory logging.info('Loading checkpoint %s...' % args.init_checkpoint_en) checkpoint_en = torch.load( os.path.join(args.init_checkpoint_en, 'checkpoint_en')) init_step_en = checkpoint_en['step'] entity_kge_model.load_state_dict(checkpoint_en['model_state_dict']) if args.do_train_en: current_learning_rate_en = checkpoint_en[ 'current_learning_rate'] warm_up_steps_en = checkpoint_en['warm_up_steps'] optimizer_en.load_state_dict( checkpoint_en['optimizer_state_dict']) else: logging.info('Ramdomly Initializing %s Model...' % args.entity_model) init_step_en = 0 step_en = init_step_en logging.info('Start Training of entity graph...') logging.info('init_step_en = %d' % init_step_en) logging.info('learning_rate_en = %f' % current_learning_rate_en) logging.info('batch_size_en = %d' % args.batch_size_en) logging.info('hidden_dim_en = %d' % args.hidden_dim_en) logging.info('gamma_en = %f' % args.gamma_en) if args.do_train_en: training_logs_en = [] for step_en in range(init_step_en, args.max_steps_en): log_en = entity_kge_model.train_step(entity_kge_model, optimizer_en, train_iterator_en, args) training_logs_en.append(log_en) if step_en >= warm_up_steps_en: current_learning_rate_en = current_learning_rate_en / 10 logging.info('Changing learning_rate_en to %f at step %d' % (current_learning_rate_en, step_en)) optimizer_en = torch.optim.Adam( filter(lambda p: p.requires_grad, entity_kge_model.parameters()), lr=current_learning_rate_en, amsgrad=True) warm_up_steps_en = warm_up_steps_en * 3 if step_en % args.save_checkpoint_steps_en == 0: save_variable_list_en = { 'step': step_en, 'current_learning_rate': current_learning_rate_en, 'warm_up_steps': warm_up_steps_en } save_model(entity_kge_model, optimizer_en, save_variable_list_en, args) if step_en % args.log_steps_en == 0: metrics_en = {} for metric_en in training_logs_en[0].keys(): metrics_en[metric_en] = sum([ log_en[metric_en] for log_en in training_logs_en ]) / len(training_logs_en) log_metrics('Training average in entity graph', step_en, metrics_en) training_logs_en = [] if args.do_valid_en and step_en % args.valid_steps_en == 0 and step_en != 0 and False: logging.info( 'Evaluating on Valid Dataset in entity graph ...') metrics_en = entity_kge_model.test_step( entity_kge_model, val_entity_triples, all_true_triples_entity, args) log_metrics('Valid', step_en, metrics_en) save_variable_list_en = { 'step': step_en, 'current_learning_rate': current_learning_rate_en, 'warm_up_steps': warm_up_steps_en } save_model(entity_kge_model, optimizer_en, save_variable_list_en, args) if args.do_valid_en and False: logging.info('Evaluating on Valid Dataset in entity graph...') metrics_en = entity_kge_model.test_step(entity_kge_model, val_entity_triples, all_true_triples_entity, args) log_metrics('Valid', step_en, metrics_en) if args.do_test_en: logging.info('Testing on Test Dataset in entity graph...') metrics_en = entity_kge_model.test_step(entity_kge_model, test_entity_triples, all_true_triples_entity, args) log_metrics('Test', step_en, metrics_en) # type else: if (not args.do_train_ty) and (not args.do_valid_ty) and ( not args.do_test_ty): raise ValueError('One of train/valid/test mode must be choosed.') if args.init_checkpoint_ty: override_config(args) elif args.data_path is None: raise ValueError( 'One of init_checkpoint/data_path must be choosed.') if args.do_train_ty and args.save_path_ty is None: raise ValueError('Where do you want to save your trained model?') if args.save_path_ty and not os.path.exists(args.save_path_ty): os.makedirs(args.save_path_ty) set_logger(args) # 获取所有type 从final_type.txt type2id = dict() with open(os.path.join(args.data_path, 'final_type_order.txt')) as fin: for line in fin: tid, type = line.strip().split('\t') type2id[type] = int(tid) ntype = len(type2id) # 获取所有type_relation 从ffinal_ty_relation_order.txt type_relation2id = dict() with open(os.path.join(args.data_path, 'ffinal_ty_relation_order.txt')) as fin: for line in fin: ty_reid, type_relation = line.strip().split('\t') type_relation2id[type_relation] = int(ty_reid) ntype_relation = len(type_relation2id) args.ntype = ntype args.ntype_re = ntype_relation logging.info('type_Model: %s' % args.type_model) logging.info('Data Path: %s' % args.data_path) logging.info('number of type: %d' % ntype) logging.info('number of type_re: %d' % ntype_relation) #加载训练三元组集合, train_type_triples = read_triple( os.path.join(args.data_path, 'train_type_Graph.txt'), type2id, type_relation2id) logging.info('#train ty_ty triples: %s' % len(train_type_triples)) #加载valid三元组集合, val_type_triples = read_triple( os.path.join(args.data_path, 'val_type_Graph.txt'), type2id, type_relation2id) logging.info('#val ty_ty triples: %s' % len(val_type_triples)) #加载test三元组集合, test_type_triples = read_triple( os.path.join(args.data_path, 'test_type_Graph.txt'), type2id, type_relation2id) logging.info('#test ty_ty triples: %s' % len(test_type_triples)) # all_triples all_true_triples_type = train_type_triples + val_type_triples + test_type_triples # type model type_kge_model = DKGE_Model( model_name=args.type_model, nnode=args.ntype, nnode_re=args.ntype_re, hidden_dim=args.hidden_dim_ty, gamma=args.gamma_ty, gamma_intra=args.gamma_intra, double_node_embedding=args.double_node_embedding_ty, double_node_re_embedding=args.double_node_re_embedding_ty) logging.info('Type Model Parameter Configuration:') for name, param in type_kge_model.named_parameters(): logging.info('Parameter %s: %s, require_grad = %s' % (name, str(param.size()), str(param.requires_grad))) if args.cuda: type_kge_model = type_kge_model.cuda() if args.do_train_ty: train_dataloader_head_ty = DataLoader( TrainDataset_ty(train_type_triples, ntype, ntype_relation, args.negative_sample_size_ty, 'head-batch'), batch_size=args.batch_size_ty, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset_ty.collate_fn) train_dataloader_tail_ty = DataLoader( TrainDataset_ty(train_type_triples, ntype, ntype_relation, args.negative_sample_size_ty, 'tail-batch'), batch_size=args.batch_size_ty, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset_ty.collate_fn) train_iterator_ty = BidirectionalOneShotIterator( train_dataloader_head_ty, train_dataloader_tail_ty) # Set training configuration current_learning_rate_ty = args.learning_rate_ty optimizer_ty = torch.optim.Adam(filter( lambda p: p.requires_grad, type_kge_model.parameters()), lr=current_learning_rate_ty, amsgrad=True) if args.warm_up_steps_ty: warm_up_steps_ty = args.warm_up_steps_ty else: warm_up_steps_ty = args.max_steps_ty // 2 if args.init_checkpoint_ty: # Restore model from checkpoint directory logging.info('Loading checkpoint %s...' % args.init_checkpoint_ty) checkpoint_ty = torch.load( os.path.join(args.init_checkpoint_ty, 'checkpoint_ty')) init_step_ty = checkpoint_ty['step'] type_kge_model.load_state_dict(checkpoint_ty['model_state_dict']) if args.do_train_ty: current_learning_rate_ty = checkpoint_ty[ 'current_learning_rate'] warm_up_steps = checkpoint_ty['warm_up_steps'] optimizer_ty.load_state_dict( checkpoint_ty['optimizer_state_dict']) else: logging.info('Ramdomly Initializing %s Model...' % args.type_model) init_step_ty = 0 step_ty = init_step_ty logging.info('Start Training of type graph...') logging.info('init_step_ty = %d' % init_step_ty) logging.info('learning_rate_ty = %f' % current_learning_rate_ty) logging.info('batch_size_ty = %d' % args.batch_size_ty) logging.info('hidden_dim_ty = %d' % args.hidden_dim_ty) logging.info('gamma_ty = %f' % args.gamma_ty) if args.do_train_ty: training_logs_ty = [] for step_ty in range(init_step_ty, args.max_steps_ty): log_ty = type_kge_model.train_step(type_kge_model, optimizer_ty, train_iterator_ty, args) training_logs_ty.append(log_ty) if step_ty >= warm_up_steps_ty: current_learning_rate_ty = current_learning_rate_ty / 10 logging.info('Changing learning_rate_ty to %f at step %d' % (current_learning_rate_ty, step_ty)) optimizer_ty = torch.optim.Adam( filter(lambda p: p.requires_grad, type_kge_model.parameters()), lr=current_learning_rate_ty, amsgrad=True) warm_up_steps_ty = warm_up_steps_ty * 3 if step_ty % args.save_checkpoint_steps_ty == 0: save_variable_list_ty = { 'step': step_ty, 'current_learning_rate': current_learning_rate_ty, 'warm_up_steps': warm_up_steps_ty } save_model(type_kge_model, optimizer_ty, save_variable_list_ty, args) if step_ty % args.log_steps_ty == 0: metrics_ty = {} for metric_ty in training_logs_ty[0].keys(): metrics_ty[metric_ty] = sum([ log_ty[metric_ty] for log_ty in training_logs_ty ]) / len(training_logs_ty) log_metrics('Training average in type graph', step_ty, metrics_ty) training_logs_ty = [] if False and args.do_valid_ty and step_ty % args.valid_steps_ty == 0: logging.info( 'Evaluating on Valid Dataset in type graph ...') metrics_ty = type_kge_model.test_step( type_kge_model, val_type_triples, all_true_triples_type, args) log_metrics('Valid', step_ty, metrics_ty) save_variable_list_ty = { 'step': step_ty, 'current_learning_rate': current_learning_rate_ty, 'warm_up_steps': warm_up_steps_ty } save_model(type_kge_model, optimizer_ty, save_variable_list_ty, args) if args.do_valid_ty and False: logging.info('Evaluating on Valid Dataset in type graph...') metrics_ty = type_kge_model.test_step(type_kge_model, val_type_triples, all_true_triples_type, args) log_metrics('Valid', step_ty, metrics_ty) if args.do_test_ty: logging.info('Testing on Test Dataset in type graph...') metrics_ty = type_kge_model.test_step(type_kge_model, test_type_triples, all_true_triples_type, args) log_metrics('Test', step_ty, metrics_ty)
def main(args): if (not args.do_train) and (not args.do_valid) and (not args.do_test): raise ValueError('one of train/val/test mode must be choosed.') if args.init_checkpoint: override_config(args) elif args.data_path is None: raise ValueError('one of init_checkpoint/data_path must be choosed.') if args.do_train and args.save_path is None: # create default save directory dt = datetime.datetime.now() args.save_path = os.path.join( os.environ['LOG_DIR'], args.data_path.split('/')[-1], args.model, datetime.datetime.now().strftime('%m%d%H%M%S')) # raise ValueError('Where do you want to save your trained model?') if args.save_path and not os.path.exists(args.save_path): os.makedirs(args.save_path) # Write logs to checkpoint and console set_logger(args) writer = SummaryWriter(log_dir=args.save_path) with open(os.path.join(args.data_path, 'entities.dict')) as fin: entity2id = dict() for line in fin: eid, entity = line.strip().split('\t') entity2id[entity] = int(eid) with open(os.path.join(args.data_path, 'relations.dict')) as fin: relation2id = dict() for line in fin: rid, relation = line.strip().split('\t') relation2id[relation] = int(rid) # Read regions for Countries S* datasets if args.countries: regions = list() with open(os.path.join(args.data_path, 'regions.list')) as fin: for line in fin: region = line.strip() regions.append(entity2id[region]) args.regions = regions nentity = len(entity2id) nrelation = len(relation2id) args.nentity = nentity args.nrelation = nrelation logging.info('Model: %s' % args.model) logging.info('Data Path: %s' % args.data_path) logging.info('Save Path: {}'.format(args.save_path)) logging.info('#entity: %d' % nentity) logging.info('#relation: %d' % nrelation) train_triples = read_triple(os.path.join(args.data_path, 'train.txt'), entity2id, relation2id) logging.info('#train: %d' % len(train_triples)) valid_triples = read_triple(os.path.join(args.data_path, 'valid.txt'), entity2id, relation2id) logging.info('#valid: %d' % len(valid_triples)) test_triples = read_triple(os.path.join(args.data_path, 'test.txt'), entity2id, relation2id) logging.info('#test: %d' % len(test_triples)) # All true triples all_true_triples = train_triples + valid_triples + test_triples if args.model in EUC_MODELS: ModelClass = EKGEModel elif args.model in HYP_MODELS: ModelClass = HKGEModel elif args.model in ONE_2_MANY_E_MODELS: ModelClass = O2MEKGEModel else: raise ValueError('model %s not supported' % args.model) if ModelClass != O2MEKGEModel: kge_model = ModelClass( model_name=args.model, nentity=nentity, nrelation=nrelation, hidden_dim=args.hidden_dim, gamma=args.gamma, p_norm=args.p_norm, dropout=args.dropout, entity_embedding_multiple=args.entity_embedding_multiple, relation_embedding_multiple=args.relation_embedding_multiple) else: kge_model = ModelClass( model_name=args.model, nentity=nentity, nrelation=nrelation, hidden_dim=args.hidden_dim, gamma=args.gamma, p_norm=args.p_norm, dropout=args.dropout, entity_embedding_multiple=args.entity_embedding_multiple, relation_embedding_multiple=args.relation_embedding_multiple, nsiblings=args.nsib, rho=args.rho) logging.info('Model Parameter Configuration:') for name, param in kge_model.named_parameters(): logging.info('Parameter %s: %s, require_grad = %s' % (name, str(param.size()), str(param.requires_grad))) if args.cuda: kge_model = kge_model.cuda() if args.do_train: # Set training dataloader iterator train_dataloader_head = DataLoader( TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, 'head-batch'), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn) train_dataloader_tail = DataLoader( TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, 'tail-batch'), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn) train_iterator = BidirectionalOneShotIterator(train_dataloader_head, train_dataloader_tail) # Set training configuration current_learning_rate = args.learning_rate optimizer = init_optimizer(kge_model, current_learning_rate) if args.warm_up_steps: warm_up_steps = args.warm_up_steps else: warm_up_steps = args.max_steps // 2 if args.init_checkpoint: # Restore model from checkpoint directory logging.info('Loading checkpoint %s...' % args.init_checkpoint) checkpoint = torch.load( os.path.join(args.init_checkpoint, 'checkpoint')) init_step = checkpoint['step'] kge_model.load_state_dict(checkpoint['model_state_dict']) if args.do_train: current_learning_rate = checkpoint['current_learning_rate'] warm_up_steps = checkpoint['warm_up_steps'] optimizer.load_state_dict(checkpoint['optimizer_state_dict']) else: logging.info('Ramdomly Initializing %s Model...' % args.model) init_step = 0 step = init_step if args.do_train: logging.info('Start Training...') logging.info('init_step = %d' % init_step) logging.info('hidden_dim = %d' % args.hidden_dim) logging.info('learning_rate = %d' % current_learning_rate) logging.info('batch_size = %d' % args.batch_size) logging.info('negative_adversarial_sampling = %d' % args.negative_adversarial_sampling) logging.info('gamma = %f' % args.gamma) logging.info('dropout = %f' % args.dropout) if args.negative_adversarial_sampling: logging.info('adversarial_temperature = %f' % args.adversarial_temperature) # Set valid dataloader as it would be evaluated during training training_logs = [] # Training Loop for step in range(init_step, args.max_steps): log = kge_model.train_step(kge_model, optimizer, train_iterator, args) training_logs.append(log) write_metrics(writer, step, log, split='train') write_metrics(writer, step, {'current_learning_rate': current_learning_rate}) if step >= warm_up_steps: current_learning_rate = current_learning_rate / 10 logging.info('Change learning_rate to %f at step %d' % (current_learning_rate, step)) optimizer = init_optimizer(kge_model, current_learning_rate) warm_up_steps = warm_up_steps * 3 if step % args.save_checkpoint_steps == 0: save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, 'warm_up_steps': warm_up_steps } save_model(kge_model, optimizer, save_variable_list, args) if step % args.log_steps == 0: metrics = {} for metric in training_logs[0].keys(): metrics[metric] = sum( [log[metric] for log in training_logs]) / len(training_logs) log_metrics('Training average', step, metrics) write_metrics(writer, step, metrics, split='train') training_logs = [] if args.do_valid and step % args.valid_steps == 0: logging.info('Evaluating on Valid Dataset...') metrics = kge_model.test_step(kge_model, valid_triples, all_true_triples, args) log_metrics('Valid', step, metrics) write_metrics(writer, step, metrics, split='valid') save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, 'warm_up_steps': warm_up_steps } save_model(kge_model, optimizer, save_variable_list, args) if args.do_valid: logging.info('Evaluating on Valid Dataset...') metrics = kge_model.test_step(kge_model, valid_triples, all_true_triples, args) log_metrics('Valid', step, metrics) if args.do_test: logging.info('Evaluating on Test Dataset...') metrics = kge_model.test_step(kge_model, test_triples, all_true_triples, args) log_metrics('Test', step, metrics) if args.evaluate_train: logging.info('Evaluating on Training Dataset...') metrics = kge_model.test_step(kge_model, train_triples, all_true_triples, args) log_metrics('Test', step, metrics)
def main(arg): with open( r'C:\Users\pc\Desktop\编程\KnowledgeGraphEmbedding-master\data\FB15k\entities.dict' ) as fin: entity2id = dict() for line in fin: eid, entity = line.strip().split('\t') entity2id[entity] = int(eid) with open( r'C:\Users\pc\Desktop\编程\KnowledgeGraphEmbedding-master\data\FB15k\relations.dict' ) as fin: relation2id = dict() for line in fin: rid, relation = line.strip().split('\t') relation2id[relation] = int(rid) nentity = len(entity2id) nrelation = len(relation2id) arg.nentity = nentity arg.nrelation = nrelation logging.info('Model:%s' % arg.MODEL) logging.info('Data Path:%s' % arg.DATA_PATH) logging.info('entity:%d' % arg.nentity) logging.info('relation:%d' % arg.nrelation) #extract data from file train_triples = read_triple(os.path.join(arg.DATA_PATH, 'train.txt'), entity2id, relation2id) logging.info('#train:%d' % len(train_triples)) valid_triples = read_triple(os.path.join(arg.DATA_PATH, 'valid.txt'), entity2id, relation2id) logging.info('#valid:%d' % len(valid_triples)) test_triples = read_triple(os.path.join(arg.DATA_PATH, 'test.txt'), entity2id, relation2id) logging.info('#test:%d' % len(test_triples)) #all true triples all_true_triples = train_triples + valid_triples + test_triples #construct model kge_model = KGEModel( model_name=arg.MODEL, nentity=arg.nentity, nrelation=arg.nrelation, hidden_dim=arg.HIDDEN_DIM, gamma=arg.gamma, double_entity_embedding=arg.double_entity_embedding, double_relation_embedding=arg.double_relation_embedding) #print model para configuration logging.info('Model Parameter Configuration') for name, para in kge_model.named_parameters(): #print(name,para.size(),para.requires_grad logging.info('Parameter %s:%s,require_grad=%s' % (name, str(para.size()), str(para.requires_grad))) #do train train_dataloader_head = DataLoader(TrainDataset(train_triples, nentity, nrelation, arg.negative_sample_size, 'head-bath'), batch_size=arg.BATCH_SIZE, shuffle=True, num_workers=max(1, arg.cpu_num // 2), collate_fn=TrainDataset.collate_fn) train_dataloader_tail = DataLoader(TrainDataset(train_triples, nentity, nrelation, arg.negative_sample_size, 'tail-batch'), batch_size=arg.BATCH_SIZE, shuffle=True, num_workers=max(1, arg.cpu_num // 2), collate_fn=TrainDataset.collate_fn) train_iterator = BidirectionalOneShotIterator(train_dataloader_head, train_dataloader_tail) #set train configuration current_learning_rate = arg.LR optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, kge_model.parameters()), lr=current_learning_rate) warm_up_steps = arg.warm_up_steps if arg.warm_up_steps else arg.max_steps // 2 init_step = 0 step = init_step logging.info('Start Training...') logging.info('init_step = %d' % init_step) logging.info('learning_rate = %d' % current_learning_rate) logging.info('batch_size = %d' % arg.BATCH_SIZE) #logging.info('negative_adversarial_sampling = %d' % arg.negative_adversarial_sampling']) logging.info('hidden_dim = %d' % arg.HIDDEN_DIM) logging.info('gamma = %f' % arg.gamma) #logging.info('negative_adversarial_sampling = %s' % str(arg.negative_adversarial_sampling'])) #start training training_logs = [] for step in range(init_step, arg.max_steps): log = kge_model.train_step(kge_model, optimizer, train_iterator, arg) training_logs.append(log) #update warm-up-step if step >= warm_up_steps: #大于warm_up_steps后学习率变为原来的1/10 current_learning_rate = current_learning_rate / 10 logging.info('Change learning_rate to %f at step %d' % (current_learning_rate, step)) optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, kge_model.parameters()), lr=current_learning_rate #更新优化器里的学习率 ) warm_up_steps = warm_up_steps * 3 #更新warm_up_steps #save model if step % arg.save_checkpoint_steps == 0: save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, 'warm_up_steps': warm_up_steps } save_model(kge_model, optimizer, save_variable_list, arg) #save after last time save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, 'warm_up_steps': warm_up_steps } save_model(kge_model, optimizer, save_variable_list, args)
def main(args): if (not args.do_train) and (not args.do_valid) and (not args.do_test) and ( not args.evaluate_train): raise ValueError('one of train/val/test mode must be choosed.') if args.init_checkpoint: override_config(args) args.save_path = 'log/%s/%s/%s-%s/%s' % ( args.dataset, args.model, args.hidden_dim, args.gamma, time.time()) if args.save_path == None else args.save_path writer = SummaryWriter(args.save_path) # Write logs to checkpoint and console set_logger(args) dataset = LinkPropPredDataset(name='ogbl-biokg') split_edge = dataset.get_edge_split() train_triples, valid_triples, test_triples = split_edge[ "train"], split_edge["valid"], split_edge["test"] nrelation = int(max(train_triples['relation'])) + 1 entity_dict = dict() cur_idx = 0 for key in dataset[0]['num_nodes_dict']: entity_dict[key] = (cur_idx, cur_idx + dataset[0]['num_nodes_dict'][key]) cur_idx += dataset[0]['num_nodes_dict'][key] nentity = sum(dataset[0]['num_nodes_dict'].values()) evaluator = Evaluator(name=args.dataset) args.nentity = nentity args.nrelation = nrelation logging.info('Model: %s' % args.model) logging.info('Dataset: %s' % args.dataset) logging.info('#entity: %d' % nentity) logging.info('#relation: %d' % nrelation) # train_triples = split_dict['train'] logging.info('#train: %d' % len(train_triples['head'])) # valid_triples = split_dict['valid'] logging.info('#valid: %d' % len(valid_triples['head'])) # test_triples = split_dict['test'] logging.info('#test: %d' % len(test_triples['head'])) train_count, train_true_head, train_true_tail = defaultdict( lambda: 4), defaultdict(list), defaultdict(list) for i in tqdm(range(len(train_triples['head']))): head, relation, tail = train_triples['head'][i], train_triples[ 'relation'][i], train_triples['tail'][i] head_type, tail_type = train_triples['head_type'][i], train_triples[ 'tail_type'][i] train_count[(head, relation, head_type)] += 1 train_count[(tail, -relation - 1, tail_type)] += 1 train_true_head[(relation, tail)].append(head) train_true_tail[(head, relation)].append(tail) kge_model = KGEModel( model_name=args.model, nentity=nentity, nrelation=nrelation, hidden_dim=args.hidden_dim, gamma=args.gamma, double_entity_embedding=args.double_entity_embedding, double_relation_embedding=args.double_relation_embedding, evaluator=evaluator) logging.info('Model Parameter Configuration:') for name, param in kge_model.named_parameters(): logging.info('Parameter %s: %s, require_grad = %s' % (name, str(param.size()), str(param.requires_grad))) if args.cuda: kge_model = kge_model.cuda() if args.init_checkpoint: # Restore model from checkpoint directory logging.info('Loading checkpoint %s...' % args.init_checkpoint) checkpoint = torch.load( os.path.join(args.init_checkpoint, 'checkpoint')) entity_dict = checkpoint['entity_dict'] if args.do_train: # Set training dataloader iterator train_dataloader_head = DataLoader( TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, 'head-batch', train_count, train_true_head, train_true_tail, entity_dict), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn) train_dataloader_tail = DataLoader( TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, 'tail-batch', train_count, train_true_head, train_true_tail, entity_dict), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn) train_iterator = BidirectionalOneShotIterator(train_dataloader_head, train_dataloader_tail) # Set training configuration current_learning_rate = args.learning_rate optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, kge_model.parameters()), lr=current_learning_rate) if args.warm_up_steps: warm_up_steps = args.warm_up_steps else: warm_up_steps = args.max_steps // 2 if args.init_checkpoint: # Restore model from checkpoint directory # logging.info('Loading checkpoint %s...' % args.init_checkpoint) # checkpoint = torch.load(os.path.join(args.init_checkpoint, 'checkpoint')) init_step = checkpoint['step'] kge_model.load_state_dict(checkpoint['model_state_dict']) # entity_dict = checkpoint['entity_dict'] if args.do_train: current_learning_rate = checkpoint['current_learning_rate'] warm_up_steps = checkpoint['warm_up_steps'] optimizer.load_state_dict(checkpoint['optimizer_state_dict']) else: logging.info('Ramdomly Initializing %s Model...' % args.model) init_step = 0 step = init_step logging.info('Start Training...') logging.info('init_step = %d' % init_step) logging.info('batch_size = %d' % args.batch_size) logging.info('negative_adversarial_sampling = %d' % args.negative_adversarial_sampling) logging.info('hidden_dim = %d' % args.hidden_dim) logging.info('gamma = %f' % args.gamma) logging.info('negative_adversarial_sampling = %s' % str(args.negative_adversarial_sampling)) if args.negative_adversarial_sampling: logging.info('adversarial_temperature = %f' % args.adversarial_temperature) # Set valid dataloader as it would be evaluated during training if args.do_train: logging.info('learning_rate = %d' % current_learning_rate) training_logs = [] #Training Loop for step in range(init_step, args.max_steps): log = kge_model.train_step(kge_model, optimizer, train_iterator, args) training_logs.append(log) if step >= warm_up_steps: current_learning_rate = current_learning_rate / 10 logging.info('Change learning_rate to %f at step %d' % (current_learning_rate, step)) optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, kge_model.parameters()), lr=current_learning_rate) warm_up_steps = warm_up_steps * 3 if step % args.save_checkpoint_steps == 0 and step > 0: # ~ 41 seconds/saving save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, 'warm_up_steps': warm_up_steps, 'entity_dict': entity_dict } save_model(kge_model, optimizer, save_variable_list, args) if step % args.log_steps == 0: metrics = {} for metric in training_logs[0].keys(): metrics[metric] = sum( [log[metric] for log in training_logs]) / len(training_logs) log_metrics('Train', step, metrics, writer) training_logs = [] if args.do_valid and step % args.valid_steps == 0 and step > 0: logging.info('Evaluating on Valid Dataset...') metrics = kge_model.test_step(kge_model, valid_triples, args, entity_dict) log_metrics('Valid', step, metrics, writer) save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, 'warm_up_steps': warm_up_steps } save_model(kge_model, optimizer, save_variable_list, args) if args.do_valid: logging.info('Evaluating on Valid Dataset...') metrics = kge_model.test_step(kge_model, valid_triples, args, entity_dict) log_metrics('Valid', step, metrics, writer) if args.do_test: logging.info('Evaluating on Test Dataset...') metrics = kge_model.test_step(kge_model, test_triples, args, entity_dict) log_metrics('Test', step, metrics, writer) if args.evaluate_train: logging.info('Evaluating on Training Dataset...') small_train_triples = {} indices = np.random.choice(len(train_triples['head']), args.ntriples_eval_train, replace=False) for i in train_triples: if 'type' in i: small_train_triples[i] = [train_triples[i][x] for x in indices] else: small_train_triples[i] = train_triples[i][indices] metrics = kge_model.test_step(kge_model, small_train_triples, args, entity_dict, random_sampling=True) log_metrics('Train', step, metrics, writer)
def main(args): if (not args.do_train) and (not args.do_valid) and (not args.do_test): raise ValueError('one of train/val/test mode must be choosed.') if args.init_checkpoint: override_config(args) elif args.data_path is None: raise ValueError('one of init_checkpoint/data_path must be choosed.') if args.do_train and args.save_path is None: raise ValueError('Where do you want to save your trained model?') if args.save_path and not os.path.exists(args.save_path): os.makedirs(args.save_path) # Write logs to checkpoint and console set_logger(args) with open(os.path.join(args.data_path, 'entities.dict')) as fin: entity2id = dict() id2entity = dict() for line in fin: eid, entity = line.strip().split('\t') entity2id[entity] = int(eid) id2entity[int(eid)] = entity with open(os.path.join(args.data_path, 'relations.dict')) as fin: relation2id = dict() id2relation = dict() for line in fin: rid, relation = line.strip().split('\t') relation2id[relation] = int(rid) id2relation[int(rid)] = relation # Read regions for Countries S* datasets if args.countries: regions = list() with open(os.path.join(args.data_path, 'regions.list')) as fin: for line in fin: region = line.strip() regions.append(entity2id[region]) args.regions = regions nentity = len(entity2id) nrelation = len(relation2id) args.nentity = nentity args.nrelation = nrelation logging.info('Model: %s' % args.model) logging.info('Data Path: %s' % args.data_path) logging.info('#entity: %d' % nentity) logging.info('#relation: %d' % nrelation) # -------------------------------------------------- # Comments by Meng: # During training, pLogicNet will augment the training triplets, # so here we load both the augmented triplets (train.txt) for training and # the original triplets (train_kge.txt) for evaluation. # Also, the hidden triplets (hidden.txt) are also loaded for annotation. # -------------------------------------------------- train_triples = read_triple(os.path.join(args.workspace_path, 'train_kge.txt'), entity2id, relation2id) logging.info('#train: %d' % len(train_triples)) train_original_triples = read_triple(os.path.join(args.data_path, 'train.txt'), entity2id, relation2id) logging.info('#train original: %d' % len(train_original_triples)) #valid_triples = read_triple(os.path.join(args.data_path, 'valid.txt'), entity2id, relation2id) #logging.info('#valid: %d' % len(valid_triples)) test_triples = read_triple(os.path.join(args.data_path, 'test.txt'), entity2id, relation2id) logging.info('#test: %d' % len(test_triples)) hidden_triples = read_triple(os.path.join(args.workspace_path, 'hidden.txt'), entity2id, relation2id) logging.info('#hidden: %d' % len(hidden_triples)) #All true triples all_true_triples = train_original_triples + test_triples kge_model = KGEModel( model_name=args.model, nentity=nentity, nrelation=nrelation, hidden_dim=args.hidden_dim, gamma=args.gamma, double_entity_embedding=args.double_entity_embedding, double_relation_embedding=args.double_relation_embedding ) logging.info('Model Parameter Configuration:') for name, param in kge_model.named_parameters(): logging.info('Parameter %s: %s, require_grad = %s' % (name, str(param.size()), str(param.requires_grad))) if args.cuda: kge_model = kge_model.cuda() if args.do_train: # Set training dataloader iterator train_dataloader_head = DataLoader( TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, 'head-batch'), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num//2), collate_fn=TrainDataset.collate_fn ) train_dataloader_tail = DataLoader( TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, 'tail-batch'), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num//2), collate_fn=TrainDataset.collate_fn ) train_iterator = BidirectionalOneShotIterator(train_dataloader_head, train_dataloader_tail) # Set training configuration current_learning_rate = args.learning_rate optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, kge_model.parameters()), lr=current_learning_rate ) if args.warm_up_steps: warm_up_steps = args.warm_up_steps else: warm_up_steps = args.max_steps // 2 if args.init_checkpoint: # Restore model from checkpoint directory logging.info('Loading checkpoint %s...' % args.init_checkpoint) checkpoint = torch.load(os.path.join(args.init_checkpoint, 'checkpoint')) init_step = checkpoint['step'] kge_model.load_state_dict(checkpoint['model_state_dict']) if args.do_train: current_learning_rate = checkpoint['current_learning_rate'] warm_up_steps = checkpoint['warm_up_steps'] optimizer.load_state_dict(checkpoint['optimizer_state_dict']) else: logging.info('Ramdomly Initializing %s Model...' % args.model) init_step = 0 step = init_step logging.info('Start Training...') logging.info('init_step = %d' % init_step) logging.info('learning_rate = %d' % current_learning_rate) logging.info('batch_size = %d' % args.batch_size) logging.info('negative_adversarial_sampling = %d' % args.negative_adversarial_sampling) logging.info('hidden_dim = %d' % args.hidden_dim) logging.info('gamma = %f' % args.gamma) logging.info('negative_adversarial_sampling = %s' % str(args.negative_adversarial_sampling)) if args.negative_adversarial_sampling: logging.info('adversarial_temperature = %f' % args.adversarial_temperature) if args.record: local_path = args.workspace_path ensure_dir(local_path) opt = vars(args) with open(local_path + '/opt.txt', 'w') as fo: for key, val in opt.items(): fo.write('{} {}\n'.format(key, val)) # Set valid dataloader as it would be evaluated during training if args.do_train: training_logs = [] #Training Loop for step in range(init_step, args.max_steps): log = kge_model.train_step(kge_model, optimizer, train_iterator, args) training_logs.append(log) if step >= warm_up_steps: current_learning_rate = current_learning_rate / 10 logging.info('Change learning_rate to %f at step %d' % (current_learning_rate, step)) optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, kge_model.parameters()), lr=current_learning_rate ) warm_up_steps = warm_up_steps * 3 if step % args.save_checkpoint_steps == 0: save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, 'warm_up_steps': warm_up_steps } save_model(kge_model, optimizer, save_variable_list, args) if step % args.log_steps == 0: metrics = {} for metric in training_logs[0].keys(): metrics[metric] = sum([log[metric] for log in training_logs])/len(training_logs) log_metrics('Training average', step, metrics) training_logs = [] if args.do_valid and (step + 1) % args.valid_steps == 0: logging.info('Evaluating on Valid Dataset...') metrics, preds = kge_model.test_step(kge_model, valid_triples, all_true_triples, args) log_metrics('Valid', step, metrics) save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, 'warm_up_steps': warm_up_steps } save_model(kge_model, optimizer, save_variable_list, args) if args.do_valid: logging.info('Evaluating on Valid Dataset...') metrics, preds = kge_model.test_step(kge_model, valid_triples, all_true_triples, args) log_metrics('Valid', step, metrics) # -------------------------------------------------- # Comments by Meng: # Save the prediction results of KGE on validation set. # -------------------------------------------------- if args.record: # Save the final results with open(local_path + '/result_kge_valid.txt', 'w') as fo: for metric in metrics: fo.write('{} : {}\n'.format(metric, metrics[metric])) # Save the predictions on test data with open(local_path + '/pred_kge_valid.txt', 'w') as fo: for h, r, t, f, rk, l in preds: fo.write('{}\t{}\t{}\t{}\t{}\n'.format(id2entity[h], id2relation[r], id2entity[t], f, rk)) for e, val in l: fo.write('{}:{:.4f} '.format(id2entity[e], val)) fo.write('\n') if args.do_test: logging.info('Evaluating on Test Dataset...') metrics, preds = kge_model.test_step(kge_model, test_triples, all_true_triples, args) log_metrics('Test', step, metrics) # -------------------------------------------------- # Comments by Meng: # Save the prediction results of KGE on test set. # -------------------------------------------------- if args.record: # Save the final results with open(local_path + '/result_kge.txt', 'w') as fo: for metric in metrics: fo.write('{} : {}\n'.format(metric, metrics[metric])) # Save the predictions on test data with open(local_path + '/pred_kge.txt', 'w') as fo: for h, r, t, f, rk, l in preds: fo.write('{}\t{}\t{}\t{}\t{}\n'.format(id2entity[h], id2relation[r], id2entity[t], f, rk)) for e, val in l: fo.write('{}:{:.4f} '.format(id2entity[e], val)) fo.write('\n') # -------------------------------------------------- # Comments by Meng: # Save the annotations on hidden triplets. # -------------------------------------------------- if args.record: # Annotate hidden triplets scores = kge_model.infer_step(kge_model, hidden_triples, args) with open(local_path + '/annotation.txt', 'w') as fo: for (h, r, t), s in zip(hidden_triples, scores): fo.write('{}\t{}\t{}\t{}\n'.format(id2entity[h], id2relation[r], id2entity[t], s)) if args.evaluate_train: logging.info('Evaluating on Training Dataset...') metrics, preds = kge_model.test_step(kge_model, train_triples, all_true_triples, args) log_metrics('Test', step, metrics)
def main(args): if (not args.do_train) and (not args.do_valid) and (not args.do_test): raise ValueError('one of train/val/test mode must be choosed.') if args.init_checkpoint: override_config(args) elif args.data_path is None: raise ValueError('one of init_checkpoint/data_path must be choosed.') if args.do_train and args.save_path is None: raise ValueError('Where do you want to save your trained model?') if args.save_path and not os.path.exists(args.save_path): os.makedirs(args.save_path) # Write logs to checkpoint and console set_logger(args) with open(os.path.join(args.data_path, 'entities.dict')) as fin: entity2id = dict() for line in fin: eid, entity = line.strip().split('\t') entity2id[entity] = int(eid) with open(os.path.join(args.data_path, 'relations.dict')) as fin: relation2id = dict() for line in fin: rid, relation = line.strip().split('\t') relation2id[relation] = int(rid) nentity = len(entity2id) nrelation = len(relation2id) args.nentity = nentity args.nrelation = nrelation logging.info('Model: %s' % args.model) logging.info('Data Path: %s' % args.data_path) logging.info('#entity: %d' % nentity) logging.info('#relation: %d' % nrelation) train_triples = read_triple(os.path.join(args.data_path, args.train_set), entity2id, relation2id) if args.fake: fake_triples = pickle.load(open(os.path.join(args.data_path, "fake%s.pkl" % args.fake), "rb")) fake = torch.LongTensor(fake_triples) train_triples += fake_triples else: fake_triples = [(0, 0, 0)] fake = torch.LongTensor(fake_triples) if args.cuda: fake = fake.cuda() logging.info('#train: %d' % len(train_triples)) valid_triples = read_triple(os.path.join(args.data_path, 'valid.txt'), entity2id, relation2id) logging.info('#valid: %d' % len(valid_triples)) test_triples = read_triple(os.path.join(args.data_path, 'test.txt'), entity2id, relation2id) logging.info('#test: %d' % len(test_triples)) all_true_triples = train_triples + valid_triples + test_triples kge_model = KGEModel( model_name=args.model, nentity=nentity, nrelation=nrelation, hidden_dim=args.hidden_dim, gamma=args.gamma, double_entity_embedding=args.double_entity_embedding, double_relation_embedding=args.double_relation_embedding ) logging.info('Model Parameter Configuration:') for name, param in kge_model.named_parameters(): logging.info('Parameter %s: %s, require_grad = %s' % (name, str(param.size()), str(param.requires_grad))) if args.cuda: kge_model = kge_model.cuda() # Set training dataloader iterator train_dataset_head = TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, 'head-batch') train_dataset_tail = TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, 'tail-batch') for triple in tqdm(train_dataset_head.triples, total=len(train_dataset_head.triples)): train_dataset_head.subsampling_weights[triple] = torch.FloatTensor([1.0]) train_dataset_tail.subsampling_weights = train_dataset_head.subsampling_weights train_dataloader_head = DataLoader( train_dataset_head, batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn ) train_dataloader_tail = DataLoader( train_dataset_tail, batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn ) train_iterator = BidirectionalOneShotIterator(train_dataloader_head, train_dataloader_tail) classifier, generator = None, None if args.method == "clf" or args.method is None: args.gen_dim = args.hidden_dim clf_triples = random.sample(train_triples, len(train_triples)//10) clf_dataset_head = TrainDataset(clf_triples, nentity, nrelation, args.negative_sample_size, 'head-batch') clf_dataset_tail = TrainDataset(clf_triples, nentity, nrelation, args.negative_sample_size, 'tail-batch') clf_dataset_head.true_head, clf_dataset_head.true_tail = train_dataset_head.true_head, train_dataset_head.true_tail clf_dataset_tail.true_head, clf_dataset_tail.true_tail = train_dataset_tail.true_head, train_dataset_tail.true_tail clf_dataset_head.subsampling_weights = train_dataset_head.subsampling_weights clf_dataset_tail.subsampling_weights = train_dataset_head.subsampling_weights clf_dataloader_head = DataLoader( clf_dataset_head, batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn ) clf_dataloader_tail = DataLoader( clf_dataset_tail, batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn ) clf_iterator = BidirectionalOneShotIterator(clf_dataloader_head, clf_dataloader_tail) gen_dataset_head = TrainDataset(clf_triples, nentity, nrelation, args.negative_sample_size, 'head-batch') gen_dataset_tail = TrainDataset(clf_triples, nentity, nrelation, args.negative_sample_size, 'tail-batch') gen_dataset_head.true_head, gen_dataset_head.true_tail = train_dataset_head.true_head, train_dataset_head.true_tail gen_dataset_tail.true_head, gen_dataset_tail.true_tail = train_dataset_tail.true_head, train_dataset_tail.true_tail gen_dataset_head.subsampling_weights = train_dataset_head.subsampling_weights gen_dataset_tail.subsampling_weights = train_dataset_head.subsampling_weights gen_dataloader_head = DataLoader( gen_dataset_head, batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn ) gen_dataloader_tail = DataLoader( gen_dataset_tail, batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num // 2), collate_fn=TrainDataset.collate_fn ) gen_iterator = BidirectionalOneShotIterator(gen_dataloader_head, gen_dataloader_tail) # if args.double_entity_embedding: # classifier = SimpleNN(input_dim=args.hidden_dim, hidden_dim=5) # generator = SimpleNN(input_dim=args.hidden_dim, hidden_dim=5) # else: classifier = SimpleNN(input_dim=args.hidden_dim, hidden_dim=5) generator = SimpleNN(input_dim=args.hidden_dim, hidden_dim=5) if args.cuda: classifier = classifier.cuda() generator = generator.cuda() clf_lr = 0.005 # if "FB15k" in args.data_path else 0.01 clf_opt = torch.optim.Adam(classifier.parameters(), lr=clf_lr) gen_opt = torch.optim.SGD(generator.parameters(), lr=0.0001) elif args.method == "KBGAN": generator = KGEModel( model_name=args.model, nentity=nentity, nrelation=nrelation, hidden_dim=args.gen_dim, gamma=args.gamma, double_entity_embedding=args.double_entity_embedding, double_relation_embedding=args.double_relation_embedding ) if args.cuda: generator = generator.cuda() # if args.gen_init is not None: # checkpoint = torch.load(os.path.join(args.gen_init, 'checkpoint')) # generator.load_state_dict(checkpoint['model_state_dict']) gen_opt = torch.optim.Adam(generator.parameters(), lr=args.learning_rate) # Set training configuration current_learning_rate = args.learning_rate optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, kge_model.parameters()), lr=current_learning_rate ) if args.warm_up_steps: warm_up_steps = args.warm_up_steps else: warm_up_steps = args.max_steps # // 2 if args.init_checkpoint: # Restore model from checkpoint directory logging.info('Loading checkpoint %s...' % args.init_checkpoint) checkpoint = torch.load(os.path.join(args.init_checkpoint, 'checkpoint')) init_step = 0 kge_model.load_state_dict(checkpoint['model_state_dict']) if args.do_train: warm_up_steps = checkpoint['warm_up_steps'] logging.info("warm_up_steps = %d" % warm_up_steps) else: current_learning_rate = args.learning_rate else: logging.info('Ramdomly Initializing %s Model...' % args.model) init_step = 0 step = init_step logging.info('Start Training...') logging.info('init_step = %d' % init_step) logging.info('learning_rate = %d' % current_learning_rate) logging.info('batch_size = %d' % args.batch_size) logging.info('negative_adversarial_sampling = %d' % args.negative_adversarial_sampling) logging.info('hidden_dim = %d' % args.hidden_dim) logging.info('gamma = %f' % args.gamma) logging.info('negative_adversarial_sampling = %s' % str(args.negative_adversarial_sampling)) if args.negative_adversarial_sampling: logging.info('adversarial_temperature = %f' % args.adversarial_temperature) # Set valid as it would be evaluated during training if args.do_train: if args.method == "clf" and args.init_checkpoint: # classifier.find_topK_triples(kge_model, classifier, train_iterator, clf_iterator, GAN_iterator) # logging.info("fake triples in classifier training %d / %d" % ( # len(set(fake_triples).intersection(set(clf_iterator.dataloader_head.dataset.triples))), # len(clf_iterator.dataloader_head.dataset.triples))) for epoch in range(1200): log = classifier.train_classifier_step(kge_model, classifier, clf_opt, clf_iterator, args, generator=None, model_name=args.model) if (epoch+1) % 200 == 0: logging.info(log) if epoch == 4000: clf_opt = torch.optim.Adam(classifier.parameters(), lr=clf_lr/10) clf_opt = torch.optim.Adam(classifier.parameters(), lr=clf_lr) training_logs = [] # Training Loop logging.info(optimizer) soft = False epoch_reward, epoch_loss, avg_reward, log = 0, 0, 0, {} for step in range(init_step, args.max_steps): if args.method == "clf" and step % 10001 == 0: if args.num == 1: soft = True elif args.num == 1000: soft = False else: soft = not soft head, relation, tail = classifier.get_embedding(kge_model, fake) if args.model == "RotatE": fake_score = classifier.forward(RotatE(head, relation, tail, "single", kge_model)) elif args.model == "DistMult": fake_score = classifier.forward(head * relation * tail) elif args.model == "TransE": fake_score = classifier.forward(head + relation - tail) all_weight = classifier.find_topK_triples(kge_model, classifier, train_iterator, clf_iterator, gen_iterator, soft=soft, model_name=args.model) logging.info("fake percent %f in %d" % (fake_score.sum().item() / all_weight, all_weight)) logging.info("fake triples in classifier training %d / %d" % ( len(set(fake_triples).intersection(set(clf_iterator.dataloader_head.dataset.triples))), len(clf_iterator.dataloader_head.dataset.triples))) epoch_reward, epoch_loss, avg_reward = 0, 0, 0 for epoch in tqdm(range(200)): classifier.train_GAN_step(kge_model, generator, classifier, gen_opt, clf_opt, gen_iterator, epoch_reward, epoch_loss, avg_reward, args, model_name=args.model) clf_train_num = 200 for epoch in range(clf_train_num): log = classifier.train_classifier_step(kge_model, classifier, clf_opt, clf_iterator, args, generator=None, model_name=args.model) if epoch % 100 == 0: logging.info(log) if step % 300 == 0 and step > 0 and args.method == "KBGAN": avg_reward = epoch_reward / batch_num epoch_reward, epoch_loss = 0, 0 logging.info('Training average reward at step %d: %f' % (step, avg_reward)) logging.info('Training average loss at step %d: %f' % (step, epoch_loss / batch_num)) if args.method == "KBGAN": epoch_reward, epoch_loss, batch_num = kge_model.train_GAN_step(generator, kge_model, gen_opt, optimizer, train_iterator, epoch_reward, epoch_loss, avg_reward, args) else: log = kge_model.train_step(kge_model, optimizer, train_iterator, args, generator=generator) training_logs.append(log) if step >= warm_up_steps: current_learning_rate = current_learning_rate / 10 logging.info('Change learning_rate to %f at step %d' % (current_learning_rate, step)) optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, kge_model.parameters()), lr=current_learning_rate ) warm_up_steps = warm_up_steps * 3 if step % args.save_checkpoint_steps == 0: save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, 'warm_up_steps': warm_up_steps } if args.method is not None: save_variable_list["confidence"] = train_iterator.dataloader_head.dataset.subsampling_weights save_model(kge_model, optimizer, save_variable_list, args, classifier=classifier, generator=generator) if step % args.log_steps == 0: metrics = {} for metric in training_logs[0].keys(): metrics[metric] = sum([log[metric] for log in training_logs]) / len(training_logs) log_metrics('Training average', step, metrics) training_logs = [] if args.do_valid and step % args.valid_steps == 0: logging.info('Evaluating on Valid Dataset...') metrics = kge_model.test_step(kge_model, valid_triples, all_true_triples, args) log_metrics('Valid', step, metrics) save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, 'warm_up_steps': warm_up_steps } if args.method is not None: save_variable_list["confidence"] = train_iterator.dataloader_head.dataset.subsampling_weights save_model(kge_model, optimizer, save_variable_list, args, classifier=classifier, generator=generator) if args.do_valid: logging.info('Evaluating on Valid Dataset...') metrics = kge_model.test_step(kge_model, valid_triples, all_true_triples, args) log_metrics('Valid', step, metrics) if args.do_test: logging.info('Evaluating on Test Dataset...') metrics = kge_model.test_step(kge_model, test_triples, all_true_triples, args) log_metrics('Test', step, metrics) if args.method is not None: classifier.find_topK_triples(kge_model, classifier, train_iterator, clf_iterator, gen_iterator, soft=True, model_name=args.model) # torch.save(train_iterator.dataloader_head.dataset.subsampling_weights, # os.path.join(args.save_path, 'weight')) true_triples = set(train_triples) - set(fake_triples) scores, label = [], [] for triple in true_triples: if not (triple == (0, 0, 0)): scores.append(train_iterator.dataloader_head.dataset.subsampling_weights[triple].item()) label.append(1) for triple in fake_triples: if not (triple == (0, 0, 0)): scores.append(train_iterator.dataloader_head.dataset.subsampling_weights[triple].item()) label.append(0) else: print("start to use sigmoid to translate distance to probability") scores, label = [], [] true_triples = set(train_triples) - set(fake_triples) i = 0 import sys while i < len(train_iterator.dataloader_head.dataset.triples): sys.stdout.write("%d in %d\r" % (i, len(train_iterator.dataloader_head.dataset.triples))) sys.stdout.flush() j = min(i + 1024, len(train_iterator.dataloader_head.dataset.triples)) sample = torch.LongTensor(train_iterator.dataloader_head.dataset.triples[i: j]).cuda() score = kge_model(sample).detach().cpu().view(-1) for x, triple in enumerate(train_iterator.dataloader_head.dataset.triples[i: j]): if triple in true_triples: label.append(1) scores.append(torch.sigmoid(score[x])) elif triple in fake_triples: label.append(0) scores.append(torch.sigmoid(score[x])) i = j del sample del score scores, label = np.array(scores), np.array(label) from sklearn.metrics import precision_score, recall_score, f1_score, roc_auc_score p = precision_score(label, scores > 0.5) r = recall_score(label, scores > 0.5) f1 = f1_score(label, scores > 0.5) auc = roc_auc_score(label, scores > 0.5) logging.info(f""" precision = {p} recall = {r} f1 score = {f1} auc score = {auc} """) p = precision_score(1 - label, scores < 0.5) r = recall_score(1 - label, scores < 0.5) f1 = f1_score(1 - label, scores < 0.5) auc = roc_auc_score(1 - label, scores < 0.5) logging.info(f""" precision = {p} recall = {r} f1 score = {f1} auc score = {auc} """) if args.evaluate_train: logging.info('Evaluating on Training Dataset...') metrics = kge_model.test_step(kge_model, train_triples, all_true_triples, args) log_metrics('Test', step, metrics)
def main(args): if args.init_checkpoint: override_config(args) elif args.data_path is None: raise ValueError('one of init_checkpoint/data_path must be choosed.') if args.save_path is None: raise ValueError('Where do you want to save your trained model?') if args.save_path and not os.path.exists(args.save_path): os.makedirs(args.save_path) # Write logs to checkpoint and console set_logger(args) with open(args.data_path) as fin: entity2id = bidict() relation2id = bidict() train_triples = [] for line in fin: _tmp = [x.strip() for x in re.split("[,\t]", line) if x.strip()][:3] if len(_tmp) < 3: continue e1, relation, e2 = tuple(_tmp) if not e1 in entity2id: entity2id[e1] = len(entity2id) if not e2 in entity2id: entity2id[e2] = len(entity2id) if not relation in relation2id: relation2id[relation] = len(relation2id) train_triples.append((entity2id[e1], relation2id[relation], entity2id[e2])) nentity = len(entity2id) nrelation = len(relation2id) args.nentity = nentity args.nrelation = nrelation logging.info('Model: %s' % args.model) logging.info('Data Path: %s' % args.data_path) logging.info('#entity: %d' % nentity) logging.info('#relation: %d' % nrelation) logging.info('#train: %d' % len(train_triples)) #All true triples all_true_triples = train_triples kge_model = KGEModel( model_name=args.model, nentity=nentity, nrelation=nrelation, hidden_dim=args.hidden_dim, gamma=args.gamma, double_entity_embedding=args.double_entity_embedding, double_relation_embedding=args.double_relation_embedding ) logging.info('Model Parameter Configuration:') for name, param in kge_model.named_parameters(): logging.info('Parameter %s: %s, require_grad = %s' % (name, str(param.size()), str(param.requires_grad))) if args.cuda: kge_model = kge_model.cuda() # Set training dataloader iterator train_dataloader_head = DataLoader( TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, 'head-batch'), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num//2), collate_fn=TrainDataset.collate_fn ) train_dataloader_tail = DataLoader( TrainDataset(train_triples, nentity, nrelation, args.negative_sample_size, 'tail-batch'), batch_size=args.batch_size, shuffle=True, num_workers=max(1, args.cpu_num//2), collate_fn=TrainDataset.collate_fn ) train_iterator = BidirectionalOneShotIterator(train_dataloader_head, train_dataloader_tail) # Set training configuration current_learning_rate = args.learning_rate optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, kge_model.parameters()), lr=current_learning_rate ) if args.warm_up_steps: warm_up_steps = args.warm_up_steps else: warm_up_steps = args.max_steps // 2 if args.init_checkpoint: # Restore model from checkpoint directory logging.info('Loading checkpoint %s...' % args.init_checkpoint) checkpoint = torch.load(os.path.join(args.init_checkpoint, 'checkpoint')) init_step = checkpoint['step'] kge_model.load_state_dict(checkpoint['model_state_dict']) current_learning_rate = checkpoint['current_learning_rate'] warm_up_steps = checkpoint['warm_up_steps'] optimizer.load_state_dict(checkpoint['optimizer_state_dict']) else: logging.info('Ramdomly Initializing %s Model...' % args.model) init_step = 0 step = init_step logging.info('Start Training...') logging.info('init_step = %d' % init_step) logging.info('learning_rate = %d' % current_learning_rate) logging.info('batch_size = %d' % args.batch_size) logging.info('negative_adversarial_sampling = %d' % args.negative_adversarial_sampling) logging.info('hidden_dim = %d' % args.hidden_dim) logging.info('gamma = %f' % args.gamma) logging.info('negative_adversarial_sampling = %s' % str(args.negative_adversarial_sampling)) if args.negative_adversarial_sampling: logging.info('adversarial_temperature = %f' % args.adversarial_temperature) # Set valid dataloader as it would be evaluated during training training_logs = [] #Training Loop for step in range(init_step, args.max_steps): log = kge_model.train_step(kge_model, optimizer, train_iterator, args) training_logs.append(log) if step >= warm_up_steps: current_learning_rate = current_learning_rate / 10 logging.info('Change learning_rate to %f at step %d' % (current_learning_rate, step)) optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, kge_model.parameters()), lr=current_learning_rate ) warm_up_steps = warm_up_steps * 3 if step % args.save_checkpoint_steps == 0: save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, 'warm_up_steps': warm_up_steps } save_model(kge_model, optimizer, save_variable_list, args, entity2id, relation2id) if step % args.log_steps == 0: metrics = {} for metric in training_logs[0].keys(): metrics[metric] = sum([log[metric] for log in training_logs])/len(training_logs) log_metrics('Training average', step, metrics) training_logs = [] save_variable_list = { 'step': step, 'current_learning_rate': current_learning_rate, 'warm_up_steps': warm_up_steps } save_model(kge_model, optimizer, save_variable_list, args, entity2id, relation2id) if args.evaluate_train: logging.info('Evaluating on Training Dataset...') metrics = kge_model.test_step(kge_model, train_triples, all_true_triples, args) log_metrics('Test', step, metrics)