def train_gat(args): # Creating the gat model here. #################################### print("Defining model") print( "\nModel type -> GAT layer with {} heads used , Initital Embeddings training" .format(args.nheads_GAT[0])) model_gat = SpKBGATModified(entity_embeddings, relation_embeddings, args.entity_out_dim, args.entity_out_dim, args.drop_GAT, args.alpha, args.nheads_GAT, args.use_simple_layer) wandb.watch(model_gat, log="all") if CUDA: model_gat.cuda() optimizer = torch.optim.Adam(model_gat.parameters(), lr=args.lr, weight_decay=args.weight_decay_gat) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=500, gamma=0.5, last_epoch=-1) torch.nn.utils.clip_grad_norm(model_gat.parameters(), 0.1) gat_loss_func = nn.MarginRankingLoss(margin=args.margin) current_batch_2hop_indices = torch.tensor([]) if (args.use_2hop): current_batch_2hop_indices = Corpus_.get_batch_nhop_neighbors_all( args, Corpus_.unique_entities_train, node_neighbors_2hop) if args.use_2hop: if CUDA: current_batch_2hop_indices = Variable( torch.LongTensor(current_batch_2hop_indices)).cuda() else: current_batch_2hop_indices = Variable( torch.LongTensor(current_batch_2hop_indices)) else: current_batch_2hop_indices = None epoch_losses = [] # losses of all epochs print("Number of epochs {}".format(args.epochs_gat)) for epoch in range(args.epochs_gat): print("\nepoch-> ", epoch) random.shuffle(Corpus_.train_triples) Corpus_.train_indices = np.array(list(Corpus_.train_triples)).astype( np.int32) model_gat.train() # getting in training mode start_time = time.time() epoch_loss = [] if len(Corpus_.train_indices) % args.batch_size_gat == 0: num_iters_per_epoch = len( Corpus_.train_indices) // args.batch_size_gat else: num_iters_per_epoch = (len(Corpus_.train_indices) // args.batch_size_gat) + 1 for iters in range(num_iters_per_epoch): start_time_iter = time.time() train_indices, train_values = Corpus_.get_iteration_batch(iters) if CUDA: train_indices = Variable( torch.LongTensor(train_indices)).cuda() train_values = Variable(torch.FloatTensor(train_values)).cuda() else: train_indices = Variable(torch.LongTensor(train_indices)) train_values = Variable(torch.FloatTensor(train_values)) # forward pass entity_embed, relation_embed = model_gat( Corpus_, Corpus_.train_adj_matrix, train_indices, current_batch_2hop_indices) optimizer.zero_grad() loss = batch_gat_loss(gat_loss_func, train_indices, entity_embed, relation_embed) loss.backward() optimizer.step() epoch_loss.append(loss.data.item()) end_time_iter = time.time() print( "Iteration-> {0} , Iteration_time-> {1:.4f} , Iteration_loss {2:.4f}" .format(iters, end_time_iter - start_time_iter, loss.data.item())) scheduler.step() print("Epoch {} , average loss {} , epoch_time {}".format( epoch, sum(epoch_loss) / len(epoch_loss), time.time() - start_time)) epoch_losses.append(sum(epoch_loss) / len(epoch_loss)) wandb.log({'epoch_loss': epoch_losses[-1]}) if (epoch + 1) % 200 == 0 or (epoch + 1) == args.epochs_gat: save_model(model_gat, args.data, epoch, args.output_folder, args.use_2hop) if (epoch + 1) == args.epochs_gat: save_final(model_gat, 'encoder', wandb.run.dir, args.use_2hop)
def train_entity_embeddings(args, word_embed_matrix, word_vocab, char_vocab): global initial_entity_emb_params global entity_embeddings model_gat = SpKBGATModified(entity_embeddings, relation_embeddings, args.entity_out_dim, args.entity_out_dim, args.drop_GAT, args.alpha, args.nheads_GAT, initial_entity_emb_params) print("Only Entity embedding model") model_entity_embedding = EntityEmbedding(\ initial_entity_emb_params['entity_embed_dim_in'], \ initial_entity_emb_params['hidden_dim_entity'], \ initial_entity_emb_params['num_encoder_layers_entity'], \ initial_entity_emb_params['is_bidirectional'], \ initial_entity_emb_params['drop_out_rate'], \ initial_entity_emb_params['hidden_dim_entity'], \ initial_entity_emb_params['entity_embed_dim_out'], \ initial_entity_emb_params['entity_conv_filter_size'], initial_entity_emb_params['word_vocab'], \ initial_entity_emb_params['word_embed_dim'], \ initial_entity_emb_params['char_embed_dim'], \ initial_entity_emb_params['word_embed_matrix'], \ initial_entity_emb_params['char_feature_size'], \ initial_entity_emb_params['conv_filter_size'], \ initial_entity_emb_params['max_word_len_entity'], \ initial_entity_emb_params['char_vocab']) if CUDA: model_gat.cuda() model_entity_embedding.cuda() if CUDA: model_gat.load_state_dict( torch.load( os.path.join(args.output_folder, 'trained_{}.pth'.format(0)))) else: model_gat.load_state_dict( torch.load(os.path.join(args.output_folder, 'trained_{}.pth'.format(0)), map_location=torch.device('cpu'))) entity_embed_gat = model_gat.entity_embeddings optimizer = torch.optim.Adam(model_gat.parameters(), lr=args.lr, weight_decay=args.weight_decay_gat) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=500, gamma=0.5, last_epoch=-1) epoch_losses = [] # losses of all epochs print("Number of epochs {}".format(args.epochs_ent_emb)) for epoch in range(args.epochs_gat): print("\nepoch-> ", epoch) random.shuffle(Corpus_.train_triples) Corpus_.train_indices = np.array(list(Corpus_.train_triples)).astype( np.int32) model_entity_embedding.train() # getting in training mode start_time = time.time() epoch_loss = [] num_emb_iters_per_epoch = int( np.ceil( len(Corpus_.unique_entities_train) / args.entities_per_batch)) random_indices = np.arange(0, len(Corpus_.unique_entities_train)) np.random.shuffle(random_indices) random_unique_entities_train = [] for idx in range(len(random_indices)): random_unique_entities_train.append( Corpus_.unique_entities_train[random_indices[idx]]) # for iters in tqdm(range(num_emb_iters_per_epoch)): start_idx = end_idx = previous_end_idx = 0 reduced_batch_size = None pbar = tqdm(total=num_emb_iters_per_epoch) iters = 0 while True: if end_idx >= len(random_unique_entities_train): break start_time_iter = time.time() try: iters += 1 start_idx = end_idx if reduced_batch_size: end_idx = min(len(random_unique_entities_train), start_idx + reduced_batch_size) else: end_idx = min(len(random_unique_entities_train), start_idx + args.entities_per_batch) current_batch_entities = random_unique_entities_train[ start_idx:end_idx] train_indices, train_values = Corpus_.get_iteration_entities_batch( current_batch_entities, args.valid_invalid_ratio_entity_embed) batch_entities_ctx_data = Corpus_.get_batch_entities_ctx_data( train_indices[:len(current_batch_entities)], args.conv_filter_size, args.max_word_len, triple=False) ctx_words = torch.from_numpy( batch_entities_ctx_data['ctx_words_list'].astype('long')) ctx_char_seq = torch.from_numpy( batch_entities_ctx_data['ctx_char_seq'].astype('long')) ctx_mask = torch.from_numpy( batch_entities_ctx_data['ctx_mask'].astype('bool')) if CUDA: train_indices = Variable( torch.LongTensor(train_indices)).cuda() train_values = Variable( torch.FloatTensor(train_values)).cuda() ctx_words = Variable(ctx_words).cuda() ctx_char_seq = Variable(ctx_char_seq).cuda() ctx_mask = Variable(ctx_mask).cuda() else: train_indices = Variable(torch.LongTensor(train_indices)) train_values = Variable(torch.FloatTensor(train_values)) ctx_words = Variable(ctx_words) ctx_char_seq = Variable(ctx_char_seq) ctx_mask = Variable(ctx_mask) entity_embeddings_batch = model_entity_embedding( ctx_words, ctx_char_seq, ctx_mask) optimizer.zero_grad() loss = batch_entity_embedding_loss(train_indices, train_values, entity_embed_gat, entity_embeddings_batch) loss.backward() optimizer.step() epoch_loss.append(loss.data.item()) reduced_batch_size = None previous_end_idx = end_idx pbar.update(1) end_time_iter = time.time() print( "Iteration-> {0}/{1} , Iteration_time-> {2:.4f} , Iteration_loss {3:.4f}" .format(iters, num_emb_iters_per_epoch, end_time_iter - start_time_iter, loss.data.item())) except RuntimeError as e: raise e scheduler.step() print("Epoch {} , average loss {} , epoch_time {}".format( epoch, sum(epoch_loss) / len(epoch_loss), time.time() - start_time)) epoch_losses.append(sum(epoch_loss) / len(epoch_loss)) save_model(model_entity_embedding, args.data, 0, args.output_folder + 'entity_embeddings/')
def train_gat(args, word_embed_matrix, word_vocab, char_vocab): # Creating the gat model here. #################################### print("Defining model") print( "\nModel type -> GAT layer with {} heads used , Initital Embeddings training" .format(args.nheads_GAT[0])) global initial_entity_emb_params global entity_embeddings model_gat = SpKBGATModified(entity_embeddings, relation_embeddings, args.entity_out_dim, args.entity_out_dim, args.drop_GAT, args.alpha, args.nheads_GAT, initial_entity_emb_params) if os.path.exists('{0}trained_{1}.pth'.format(args.output_folder, 0)): if CUDA: model_gat.load_state_dict( torch.load( os.path.join(args.output_folder, 'trained_{}.pth'.format(0)))) else: model_gat.load_state_dict( torch.load(os.path.join(args.output_folder, 'trained_{}.pth'.format(0)), map_location=torch.device('cpu'))) display_every = 100 if CUDA: model_gat.cuda() # optimizer = torch.optim.Adam( # model_gat.parameters(), lr=args.lr, weight_decay=args.weight_decay_gat) optimizer = torch.optim.SGD(model_gat.parameters(), lr=args.lr) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=500, gamma=0.5, last_epoch=-1) gat_loss_func = nn.MarginRankingLoss(margin=args.margin) epoch_losses = [] # losses of all epochs print("Number of epochs {}".format(args.epochs_gat)) for epoch in range(args.epochs_gat): print("\nepoch-> ", epoch) random.shuffle(Corpus_.train_triples) Corpus_.train_indices = np.array(list(Corpus_.train_triples)).astype( np.int32) model_gat.train() # getting in training mode start_time = time.time() epoch_loss = [] num_iters_per_epoch = int( np.ceil(len(Corpus_.unique_entities_train) / args.batch_size_gat)) random_indices = np.arange(0, len(Corpus_.unique_entities_train)) np.random.shuffle(random_indices) random_unique_entities_train = [] for idx in range(len(random_indices)): random_unique_entities_train.append( Corpus_.unique_entities_train[random_indices[idx]]) start_idx = end_idx = previous_end_idx = 0 reduced_batch_size = None iters = 0 for iters in tqdm(range(num_iters_per_epoch)): start_time_iter = time.time() start_idx = end_idx if reduced_batch_size: end_idx = min(len(random_unique_entities_train), start_idx + reduced_batch_size) else: end_idx = min(len(random_unique_entities_train), start_idx + args.entities_per_batch) batch_train_adj_matrix, current_batch_entities_set = Corpus_.get_batch_adj_data( iters, start_idx=start_idx, end_idx=end_idx, unique_entities_train=random_unique_entities_train) if batch_train_adj_matrix[0][0].shape[0] == 0: continue train_indices, train_values = Corpus_.get_iteration_triples_batch( list(current_batch_entities_set['source'])) if args.use_2hop: current_batch_2hop_indices = Corpus_.get_batch_nhop_neighbors_all( args, list(current_batch_entities_set['source']), node_neighbors_2hop) if current_batch_2hop_indices.shape[0] == 0: # current_batch_2hop_indices = np.empty((1,4)) current_batch_2hop_indices = torch.tensor([], dtype=torch.long) else: current_batch_2hop_indices = torch.tensor([], dtype=torch.long) if CUDA: train_indices = Variable( torch.LongTensor(train_indices)).cuda() train_values = Variable(torch.FloatTensor(train_values)).cuda() if args.use_2hop: current_batch_2hop_indices = Variable( torch.LongTensor(current_batch_2hop_indices)).cuda() else: train_indices = Variable(torch.LongTensor(train_indices)) train_values = Variable(torch.FloatTensor(train_values)) if args.use_2hop: current_batch_2hop_indices = Variable( torch.LongTensor(current_batch_2hop_indices)) entity_embed, relation_embed, mask = model_gat( Corpus_, torch.tensor(list(current_batch_entities_set['source'])), batch_train_adj_matrix, current_batch_2hop_indices) optimizer.zero_grad() loss = batch_gat_loss(gat_loss_func, train_indices, entity_embed, relation_embed) loss.backward() optimizer.step() epoch_loss.append(loss.data.item()) end_time_iter = time.time() print( "Iteration-> {0}/{1} , Iteration_time-> {2:.4f} , Iteration_loss {3:.4f}" .format(iters, num_iters_per_epoch, end_time_iter - start_time_iter, loss.data.item())) if iters % display_every == 0: num_samples = 1000 random_ent_indices = np.random.randint(0, entity_embed.shape[0], num_samples) sampled_entities = np.array( entity_embed[random_ent_indices].detach().cpu()) mean_vector = np.mean(sampled_entities, axis=-1) norm_entities = np.sqrt( np.sum(np.square(sampled_entities), axis=-1)) norm_mean = np.sqrt( np.sum(np.square(sampled_entities), axis=-1)) den = norm_mean * norm_entities num = np.dot(mean_vector, norm_entities.transpose()) cosine_dist = num / den mean_cosine_dist = np.mean(cosine_dist) median_cosine_dist = np.median(cosine_dist) min_norm = np.min(norm_entities) max_norm = np.max(norm_entities) print('mean_cosine_dist: ', mean_cosine_dist) print('median_cosine_dist: ', median_cosine_dist) print('min_norm: ', min_norm) print('max_norm: ', max_norm) scheduler.step() print("Epoch {} , average loss {} , epoch_time {}".format( epoch, sum(epoch_loss) / len(epoch_loss), time.time() - start_time)) epoch_losses.append(sum(epoch_loss) / len(epoch_loss)) save_model(model_gat, args.data, 0, args.output_folder)
def train_gat(args): # Creating the gat model here. #################################### print("Defining model") print( "\nModel type -> GAT layer with {} heads used , Initital Embeddings training" .format(args.nheads_GAT[0])) model_gat = SpKBGATModified(entity_embeddings, relation_embeddings, args.entity_out_dim, args.entity_out_dim, args.drop_GAT, args.alpha, args.nheads_GAT) if CUDA: model_gat.cuda() optimizer = torch.optim.Adam(model_gat.parameters(), lr=args.lr, weight_decay=args.weight_decay_gat) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=500, gamma=0.5, last_epoch=-1) gat_loss_func = nn.MarginRankingLoss(margin=args.margin) current_batch_2hop_indices = torch.tensor([]) if (args.use_2hop): current_batch_2hop_indices = Corpus_.get_batch_nhop_neighbors_all( args, Corpus_.unique_entities_train, node_neighbors_2hop) if CUDA: current_batch_2hop_indices = Variable( torch.LongTensor(current_batch_2hop_indices)).cuda() else: current_batch_2hop_indices = Variable( torch.LongTensor(current_batch_2hop_indices)) epoch_losses = [] # losses of all epochs print("Number of epochs {}".format(args.epochs_gat)) for epoch in range(args.epochs_gat): random.shuffle(Corpus_.train_triples) Corpus_.train_indices = np.array(list(Corpus_.train_triples)).astype( np.int32) model_gat.train() # getting in training mode start_time = time.time() epoch_loss = [] if len(Corpus_.train_indices) % args.batch_size_gat == 0: num_iters_per_epoch = len( Corpus_.train_indices) // args.batch_size_gat else: num_iters_per_epoch = (len(Corpus_.train_indices) // args.batch_size_gat) + 1 for iters in range(num_iters_per_epoch): start_time_iter = time.time() train_indices, train_values = Corpus_.get_iteration_batch(iters) if CUDA: train_indices = Variable( torch.LongTensor(train_indices)).cuda() train_values = Variable(torch.FloatTensor(train_values)).cuda() else: train_indices = Variable(torch.LongTensor(train_indices)) train_values = Variable(torch.FloatTensor(train_values)) # forward pass entity_embed, relation_embed = model_gat( Corpus_, Corpus_.train_adj_matrix, train_indices, current_batch_2hop_indices) optimizer.zero_grad() loss = batch_gat_loss(gat_loss_func, train_indices, entity_embed, relation_embed) loss.backward() optimizer.step() epoch_loss.append(loss.data.item()) end_time_iter = time.time() scheduler.step() if epoch % 100 == 0: print("Epoch {} , average loss {} , epoch_time {}".format( epoch, sum(epoch_loss) / len(epoch_loss), time.time() - start_time)) epoch_losses.append(sum(epoch_loss) / len(epoch_loss)) summary.add_scalar('loss/gat_loss', sum(epoch_loss) / len(epoch_loss), epoch) if (epoch + 1) % 200 == 0: save_model(model_gat, args.data, epoch, args.output_folder) now = time.localtime() f = open( (args.output_folder + "train_gat_epoch_losses_{}-{}-{}.txt").format( now.tm_year, now.tm_mon, now.tm_mday), 'w') for i in epoch_losses: f.write(str(i)) f.write('\n') f.close()
def train_gat(args): # Creating the gat model here. #################################### print("Defining model") print("\nModel type -> GAT layer with {} heads used , Initital Embeddings training".format(args.nheads_GAT[0])) ''' model_gat初始化进行修改 ''' model_gat = SpKBGATModified(entity_embeddings, relation_embeddings, nhop_embeddings, nhop_array, args.entity_out_dim, args.entity_out_dim, args.drop_GAT, args.alpha, args.nheads_GAT) if CUDA: model_gat = model_gat.cuda() optimizer = torch.optim.Adam(model_gat.parameters(), lr=args.lr, weight_decay=args.weight_decay_gat) scheduler = torch.optim.lr_scheduler.StepLR( optimizer, step_size=500, gamma=0.5, last_epoch=-1) gat_loss_func = nn.MarginRankingLoss(margin=args.margin) # train_loader = torch.utils.data.DataLoader( # current_batch_2hop_indices, sampler=train_sampler, **kwargs) epoch_losses = [] # losses of all epochs print("Number of epochs {}".format(args.epochs_gat)) """ Add by cc Add early stopping with patience if patience times in a row score_fn gives a result lower than the best result, than the training will be stopped """ counter = 0 best_score = None for epoch in range(args.epochs_gat): print("\nepoch-> ", epoch) random.shuffle(Corpus_.train_triples) Corpus_.train_indices = np.array( list(Corpus_.train_triples)).astype(np.int32) model_gat.train() # getting in training mode start_time = time.time() epoch_loss = [] if len(Corpus_.train_indices) % args.batch_size_gat == 0: num_iters_per_epoch = len(Corpus_.train_indices) // args.batch_size_gat else: num_iters_per_epoch = (len(Corpus_.train_indices) // args.batch_size_gat) + 1 for iters in range(num_iters_per_epoch): start_time_iter = time.time() train_indices, train_values = Corpus_.get_iteration_batch(iters) iter_unique_entities = get_unique_entity(train_indices) current_batch_2hop_indices = torch.tensor([]) start_2hop_time = time.time() current_batch_2hop_indices = Corpus_.get_batch_nhop_neighbors_all(args, iter_unique_entities, node_neighbors_2hop) end_2hop_time = time.time() print(end_2hop_time - start_2hop_time) if CUDA: train_indices = Variable(torch.LongTensor(train_indices)).cuda() # train_values = Variable(torch.FloatTensor(train_values)).cuda().half() else: train_indices = Variable(torch.LongTensor(train_indices)) # train_values = Variable(torch.FloatTensor(train_values)) current_batch_2hop_indices = Variable(torch.LongTensor(current_batch_2hop_indices)).cuda() # forward pass entity_embed, relation_embed = model_gat( Corpus_, Corpus_.train_adj_matrix, train_indices, current_batch_2hop_indices) optimizer.zero_grad() loss = batch_gat_loss( gat_loss_func, train_indices, entity_embed, relation_embed) loss.backward() optimizer.step() epoch_loss.append(loss.data.item()) end_time_iter = time.time() print("Iteration-> {0} , Iteration_time-> {1:.4f} , Iteration_loss {2:.4f}".format( iters, end_time_iter - start_time_iter, loss.data.item())) scheduler.step() avg_loss = sum(epoch_loss) / len(epoch_loss) print("Epoch {} , average loss {} , epoch_time {}".format( epoch, avg_loss, time.time() - start_time)) epoch_losses.append(avg_loss) """ early stopping """ if best_score is None: best_score = 99 elif avg_loss > best_score: counter += 1 if counter >= args.patience_gat: break else: best_score = avg_loss counter = 0 # plot avg_loss writer.add_scalar('WN18RR_testing [add (h,t)->r]: GAT average loss--epoch', sum(epoch_loss) / len(epoch_loss), epoch) save_model(model_gat, args.data, epoch, args.output_folder)
def train_gat(args): # Creating the gat model here. #################################### print("Defining model") print( "\nModel type -> GAT layer with {} heads used , Initital Embeddings training".format(args.nheads_GAT[0])) if args.tanh: model_gat = SpKBGATModified(entity_embeddings, relation_embeddings, args.entity_out_dim, args.entity_out_dim, args.drop_GAT, args.alpha, args.nheads_GAT, 'tanh') else: model_gat = SpKBGATModified(entity_embeddings, relation_embeddings, args.entity_out_dim, args.entity_out_dim, args.drop_GAT, args.alpha, args.nheads_GAT, 'leakyrelu') model_gat = nn.DataParallel(model_gat) if CUDA: model_gat.cuda() if args.load_gat is not None and args.epochs_gat == 0: model_gat.load_state_dict(torch.load( '{0}gat/trained_{1}.pth'.format(args.output_folder, args.epochs_gat - 1))) elif args.load_gat is not None and args.epochs_gat > 0: model_gat.load_state_dict(torch.load(args.load_gat)) optimizer = torch.optim.Adam( model_gat.parameters(), lr=args.lr, weight_decay=args.weight_decay_gat) scheduler = torch.optim.lr_scheduler.StepLR( optimizer, step_size=500, gamma=0.5, last_epoch=-1) gat_loss_func = nn.MarginRankingLoss(margin=args.margin) current_batch_2hop_indices = torch.tensor([], dtype=torch.long) if(args.use_2hop): current_batch_2hop_indices = Corpus_.get_batch_nhop_neighbors_all(args, Corpus_.unique_entities_train, node_neighbors_2hop) if CUDA: current_batch_2hop_indices = Variable( torch.LongTensor(current_batch_2hop_indices)).cuda() else: current_batch_2hop_indices = Variable( torch.LongTensor(current_batch_2hop_indices)) epoch_losses = [] # losses of all epochs print("Number of epochs {}".format(args.epochs_gat)) for epoch in range(args.epochs_gat): print("\nepoch-> ", epoch) random.shuffle(Corpus_.train_triples) Corpus_.train_indices = np.array( list(Corpus_.train_triples)).astype(np.int32) model_gat.train() # getting in training mode start_time = time.time() epoch_loss = [] if len(Corpus_.train_indices) % args.batch_size_gat == 0: num_iters_per_epoch = len( Corpus_.train_indices) // args.batch_size_gat else: num_iters_per_epoch = ( len(Corpus_.train_indices) // args.batch_size_gat) + 1 if args.debug: pass for iters in range(num_iters_per_epoch): start_time_iter = time.time() train_indices, train_values = Corpus_.get_iteration_batch(iters) if CUDA: train_indices = Variable( torch.LongTensor(train_indices)).cuda() train_values = Variable(torch.FloatTensor(train_values)).cuda() else: train_indices = Variable(torch.LongTensor(train_indices)) train_values = Variable(torch.FloatTensor(train_values)) # forward pass entity_embed, relation_embed = model_gat( Corpus_, Corpus_.train_adj_matrix, train_indices, current_batch_2hop_indices) # print('Forward pass', entity_embed.shape, relation_embed.shape) optimizer.zero_grad() loss = batch_gat_loss( gat_loss_func, train_indices, entity_embed, relation_embed) loss.backward() optimizer.step() epoch_loss.append(loss.data.item()) end_time_iter = time.time() line = "Iteration-> {0} , Iteration_time-> {1:.4f} , Iteration_loss {2:.4f}".format( iters, end_time_iter - start_time_iter, loss.data.item()) print(line) write_to_file(args, line) scheduler.step() line = "Epoch {} , average loss {} , epoch_time {}".format( epoch, sum(epoch_loss) / len(epoch_loss), time.time() - start_time) print(line) write_to_file(args, line) epoch_losses.append(sum(epoch_loss) / len(epoch_loss)) save_model(model_gat, args.data, epoch, args.output_folder + 'gat/') if CUDA: final_entity_embeddings = model_gat.module.final_entity_embeddings.cpu().detach().numpy() final_relation_embeddings = model_gat.module.final_relation_embeddings.cpu().detach().numpy() else: final_entity_embeddings = model_gat.module.final_entity_embeddings.detach().numpy() final_relation_embeddings = model_gat.module.final_relation_embeddings.detach().numpy() save_mbeddings(args, final_entity_embeddings, final_relation_embeddings)