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)
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, '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 train_evaluate_triples = sample(train_triples, 4000) kge_model = KGEModel(model_name=args.model, nentity=nentity, nrelation=nrelation, hidden_dim=args.hidden_dim, gamma=args.gamma) 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=1, 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=1, 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) 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('hidden_dim = %d' % args.hidden_dim) logging.info('gamma = %f' % args.gamma) # Set valid dataloader as it would be evaluated during training if args.do_train: logging.info('learning_rate = %f' % current_learning_rate) training_logs = [] #writer = SummaryWriter('loss') #Training Loop for step in range(init_step, args.max_steps): log, loss = kge_model.train_step(kge_model, optimizer, train_iterator, args) training_logs.append(log) 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.max_steps - 1000 and step % 100 == 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)
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)
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() id2relationship = dict() for line in fin: rid, relation = line.strip().split('\t') relation2id[relation] = int(rid) id2relationship[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 e_vocab = Dictionary(tok2ind=entity2id, ind2tok=id2entity) r_vocab = Dictionary(tok2ind=relation2id, ind2tok=id2relationship) ## TODO: add graph file graph = KB(os.path.join(args.data_path, 'train.txt'), e_vocab=e_vocab, r_vocab=r_vocab) 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)) candidate_entities = None if args.rerank_minerva: candidate_entities = defaultdict(set) with open( "/home/shdhulia//Limits-of-Path-Reasoner/outputs/FB15K-237/thisone_test/all_answers.txt" ) as candidate_file: # with open("/home/shdhulia/minerva_answers/fb.txt") as candidate_file: for line in candidate_file: pt = line.strip().split("\t") e1 = entity2id[pt[0]] r = relation2id[pt[1]] predicted_es = set( [entity2id[p] for p in pt[2:] if p in entity2id]) candidate_entities[(e1, r)] = set(predicted_es) #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', KB=graph), 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) train_iterator = OneShotIterator(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, candidate_entities, id2e=id2entity, id2rel=id2relationship) 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.do_test: logging.info('Evaluating on Test Dataset...') metrics = kge_model.test_step(kge_model, test_triples, all_true_triples, args, candidate_entities, id2e=id2entity, id2rel=id2relationship) 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)
print("start training") init_step = 1 total_steps = 500001 test_steps = 10000 last_loss = 100 score = 0 need_to_load_checkpoint = True if need_to_load_checkpoint: _, init_step, score, last_loss = load_checkpoint(model, optim, checkpoint_path) progbar = Progbar(max_step=total_steps - init_step) start_time = time.time() for step in range(init_step, total_steps): loss = model.train_step(model, optim, train_iterator, device) progbar.update(step - init_step, [ ("step", step - init_step), ("loss", loss), ("cost", round((time.time() - start_time))) ]) summary_writer.add_scalar(tag='Loss/train', scalar_value=loss, global_step=step) if step > init_step and step % test_steps == 0: print("\n属性消融实验") left_vec = t.get_vec2(model.entity_embedding, t.left_ids) right_vec = t.get_vec2(model.entity_embedding, t.right_ids) hits = t.get_hits(left_vec, right_vec) hits_left = hits["left"] hits_right = hits["right"] left_hits_10 = hits_left[2][1]