| import math |
| import multiprocessing as mp |
| import random |
| import time |
| import gc |
|
|
| import tensorflow as tf |
| import numpy as np |
| from sklearn import preprocessing |
|
|
| import openea.modules.train.batch as bat |
| from openea.modules.finding.evaluation import early_stop |
| from openea.modules.utils.util import generate_out_folder |
| from openea.modules.utils.util import load_session |
| from openea.modules.utils.util import task_divide |
| from openea.modules.base.initializers import init_embeddings, orthogonal_init |
| from openea.modules.base.losses import margin_loss |
| from openea.modules.base.optimizers import generate_optimizer |
|
|
| from utils import MyKGs, find_alignment |
| from eval import valid, test |
| from nn_search import generate_neighbours |
|
|
|
|
| def search_kg1_to_kg2_1nn_neighbor(embeds1, embeds2, ents2, mapping_mat, return_sim=False, soft_nn=10): |
| if mapping_mat is not None: |
| embeds1 = np.matmul(embeds1, mapping_mat) |
| embeds1 = preprocessing.normalize(embeds1) |
| sim_mat = np.matmul(embeds1, embeds2.T) |
| nearest_pairs = find_alignment(sim_mat, soft_nn) |
| nns = [ents2[x[0][1]] for x in nearest_pairs] |
| if return_sim: |
| sim_list = [] |
| for pair in nearest_pairs: |
| sim_list.append(sim_mat[pair[0][0], pair[0][1]]) |
| return nns, sim_list |
| return nns |
|
|
|
|
| class MTransEV2: |
|
|
| def set_kgs(self, kgs: MyKGs): |
| self.kgs = kgs |
|
|
| def set_args(self, args): |
| self.args = args |
| self.out_folder = generate_out_folder(self.args.output, self.args.training_data, self.args.dataset_division, |
| self.__class__.__name__) |
|
|
| def init(self): |
| self._define_variables() |
| self._define_embed_graph() |
| self._define_align_graph() |
|
|
| self.session = load_session() |
| tf.global_variables_initializer().run(session=self.session) |
|
|
| def __init__(self): |
|
|
| self.out_folder = None |
| self.args = None |
| self.kgs = None |
|
|
| self.session = None |
|
|
| self.seed_entities1 = None |
| self.seed_entities2 = None |
| self.neg_ts = None |
| self.neg_rs = None |
| self.neg_hs = None |
| self.pos_ts = None |
| self.pos_rs = None |
| self.pos_hs = None |
|
|
| self.rel_embeds = None |
| self.ent_embeds = None |
| self.mapping_mat = None |
| self.eye_mat = None |
|
|
| self.triple_optimizer = None |
| self.triple_loss = None |
| self.mapping_optimizer = None |
| self.mapping_loss = None |
|
|
| self.mapping_mat = None |
|
|
| self.flag1 = -1 |
| self.flag2 = -1 |
| self.early_stop = False |
|
|
| def _define_variables(self): |
| with tf.variable_scope('KG' + 'embeddings'): |
| self.ent_embeds = init_embeddings([self.kgs.entities_num, self.args.dim], 'ent_embeds', |
| self.args.init, self.args.ent_l2_norm) |
| self.rel_embeds = init_embeddings([self.kgs.relations_num, self.args.dim], 'rel_embeds', |
| self.args.init, self.args.rel_l2_norm) |
|
|
| with tf.variable_scope('alignment' + 'mapping'): |
| self.mapping_mat = orthogonal_init([self.args.dim, self.args.dim], 'mapping_matrix') |
| self.eye_mat = tf.constant(np.eye(self.args.dim), dtype=tf.float32, name='eye') |
|
|
| def _define_embed_graph(self): |
| print("build embedding learning graph...") |
| with tf.name_scope('triple_placeholder'): |
| self.pos_hs = tf.placeholder(tf.int32, shape=[None]) |
| self.pos_rs = tf.placeholder(tf.int32, shape=[None]) |
| self.pos_ts = tf.placeholder(tf.int32, shape=[None]) |
| self.neg_hs = tf.placeholder(tf.int32, shape=[None]) |
| self.neg_rs = tf.placeholder(tf.int32, shape=[None]) |
| self.neg_ts = tf.placeholder(tf.int32, shape=[None]) |
| with tf.name_scope('triple_lookup'): |
| phs = tf.nn.embedding_lookup(self.ent_embeds, self.pos_hs) |
| prs = tf.nn.embedding_lookup(self.rel_embeds, self.pos_rs) |
| pts = tf.nn.embedding_lookup(self.ent_embeds, self.pos_ts) |
| nhs = tf.nn.embedding_lookup(self.ent_embeds, self.neg_hs) |
| nrs = tf.nn.embedding_lookup(self.rel_embeds, self.neg_rs) |
| nts = tf.nn.embedding_lookup(self.ent_embeds, self.neg_ts) |
| with tf.name_scope('triple_loss'): |
| self.triple_loss = margin_loss(phs, prs, pts, nhs, nrs, nts, self.args.embed_margin, self.args.loss_norm) |
| self.triple_optimizer = generate_optimizer(self.triple_loss, self.args.learning_rate, |
| opt=self.args.optimizer) |
|
|
| def _mapping_align_loss(self, seed_embed1, seed_embed2): |
| mapped_seed_embed1 = tf.matmul(seed_embed1, self.mapping_mat) |
| mapped_seed_embed1 = tf.nn.l2_normalize(mapped_seed_embed1, 1) |
| with tf.name_scope('mapping_distance'): |
| distance = mapped_seed_embed1 - seed_embed2 |
| with tf.name_scope('mapping_loss'): |
| align_loss = tf.reduce_sum(tf.reduce_sum(tf.square(distance), axis=1)) |
| orthogonal_loss = tf.reduce_mean( |
| tf.reduce_sum(tf.pow(tf.matmul(self.mapping_mat, self.mapping_mat, transpose_b=True) - self.eye_mat, 2), 1)) |
| return align_loss + orthogonal_loss |
|
|
| def _mapping_align_marginal_loss(self, seed_embed1, seed_embed2, pos_embed1, neg_embed2): |
| mapped_seed_embed1 = tf.matmul(seed_embed1, self.mapping_mat) |
| mapped_seed_embed1 = tf.nn.l2_normalize(mapped_seed_embed1, 1) |
| with tf.name_scope('mapping_distance'): |
| distance = mapped_seed_embed1 - seed_embed2 |
| with tf.name_scope('mapping_loss'): |
| pos_score = tf.reduce_sum(tf.square(distance), axis=1) |
| align_loss = tf.reduce_sum(pos_score) |
| if self.args.mapping_margin > 0.0: |
| mapped_pos_embed1 = tf.matmul(pos_embed1, self.mapping_mat) |
| mapped_pos_embed1 = tf.nn.l2_normalize(mapped_pos_embed1, 1) |
| neg_distance = mapped_pos_embed1 - neg_embed2 |
| neg_score = tf.reduce_sum(tf.square(neg_distance), axis=1) |
| align_loss += self.args.mapping_neg_weight * \ |
| tf.reduce_sum(tf.nn.relu(tf.constant(self.args.mapping_margin) - neg_score)) |
| orthogonal_loss = tf.reduce_mean(tf.reduce_sum( |
| tf.square(tf.matmul(self.mapping_mat, self.mapping_mat, transpose_b=True) - self.eye_mat), 1)) |
| return align_loss + orthogonal_loss |
|
|
| def _define_align_graph(self): |
| print("build alignment learning graph...") |
| with tf.name_scope('seed_links_placeholder'): |
| self.seed_entities1 = tf.placeholder(tf.int32, shape=[None]) |
| self.seed_entities2 = tf.placeholder(tf.int32, shape=[None]) |
| self.pos_entities1 = tf.placeholder(tf.int32, shape=[None]) |
| self.neg_entities2 = tf.placeholder(tf.int32, shape=[None]) |
| with tf.name_scope('seed_links_lookup'): |
| seed_embed1 = tf.nn.embedding_lookup(self.ent_embeds, self.seed_entities1) |
| seed_embed2 = tf.nn.embedding_lookup(self.ent_embeds, self.seed_entities2) |
| pos_embed1 = tf.nn.embedding_lookup(self.ent_embeds, self.pos_entities1) |
| neg_embed2 = tf.nn.embedding_lookup(self.ent_embeds, self.neg_entities2) |
| with tf.name_scope('mapping_loss'): |
| self.mapping_loss = self._mapping_align_marginal_loss(seed_embed1, seed_embed2, pos_embed1, neg_embed2) |
| self.mapping_optimizer = generate_optimizer(self.mapping_loss, self.args.learning_rate, |
| opt=self.args.optimizer) |
|
|
| def _eval_valid_embeddings(self, remove_dangling=False): |
| if remove_dangling: |
| candidate_list = self.kgs.valid_entities2 + list(self.kgs.kg2.entities_set |
| - set(self.kgs.train_entities2) |
| - set(self.kgs.valid_entities2) |
| - set([x for x, _ in self.kgs.train_unlinked_entities2]) |
| - set([x for x, _ in self.kgs.valid_unlinked_entities2])) |
| else: |
| candidate_list = self.kgs.valid_entities2 + list(self.kgs.kg2.entities_set |
| - set(self.kgs.train_entities2) |
| - set(self.kgs.valid_entities2)) |
| embeds1 = tf.nn.embedding_lookup(self.ent_embeds, self.kgs.valid_entities1).eval(session=self.session) |
| embeds2 = tf.nn.embedding_lookup(self.ent_embeds, candidate_list).eval(session=self.session) |
| mapping = self.mapping_mat.eval(session=self.session) if self.mapping_mat is not None else None |
| return embeds1, embeds2, mapping |
|
|
| def _eval_test_embeddings(self, remove_dangling=False): |
| if remove_dangling: |
| candidate_list = self.kgs.test_entities2 + list(self.kgs.kg2.entities_set |
| - set(self.kgs.train_entities2) |
| - set(self.kgs.valid_entities2) |
| - set(self.kgs.test_entities2) |
| - set([x for x, _ in self.kgs.train_unlinked_entities2]) |
| - set([x for x, _ in self.kgs.valid_unlinked_entities2])) |
| else: |
| candidate_list = self.kgs.test_entities2 + list(self.kgs.kg2.entities_set |
| - set(self.kgs.train_entities2) |
| - set(self.kgs.valid_entities2) |
| - set(self.kgs.test_entities2)) |
| embeds1 = tf.nn.embedding_lookup(self.ent_embeds, self.kgs.test_entities1).eval(session=self.session) |
| embeds2 = tf.nn.embedding_lookup(self.ent_embeds, candidate_list).eval(session=self.session) |
| mapping = self.mapping_mat.eval(session=self.session) if self.mapping_mat is not None else None |
| return embeds1, embeds2, mapping |
|
|
| def valid_alignment(self, stop_metric, remove_dangling): |
| print("\nvalidating entity alignment...") |
| embeds1, embeds2, mapping = self._eval_valid_embeddings(remove_dangling=remove_dangling) |
| hits, mrr_12, sim_list = valid(embeds1, embeds2, mapping, self.args.top_k, |
| self.args.test_threads_num, metric=self.args.eval_metric, |
| normalize=self.args.eval_norm, csls_k=0, accurate=False) |
| print() |
| return hits[0] if stop_metric == 'hits1' else mrr_12 |
|
|
| def test(self): |
| embeds1, embeds2, mapping = self._eval_test_embeddings(remove_dangling=False) |
| _, _, _, sim_list = test(embeds1, embeds2, mapping, self.args.top_k, self.args.test_threads_num, |
| metric=self.args.eval_metric, normalize=self.args.eval_norm, csls_k=0, accurate=True) |
| print() |
|
|
| def eval_embeddings(self, entity_list): |
| embeds = tf.nn.embedding_lookup(self.ent_embeds, entity_list) |
| return embeds.eval(session=self.session) |
|
|
| def launch_training_1epo(self, epoch, triple_steps, steps_tasks, training_batch_queue, neighbors1, neighbors2): |
| self.launch_embed_training_1epo(epoch, triple_steps, steps_tasks, training_batch_queue, neighbors1, neighbors2) |
| self.launch_align_training_1epo(epoch, triple_steps) |
|
|
| def launch_embed_training_1epo(self, epoch, triple_steps, steps_tasks, batch_queue, neighbors1, neighbors2): |
| start = time.time() |
| for steps_task in steps_tasks: |
| mp.Process(target=bat.generate_relation_triple_batch_queue, |
| args=(self.kgs.kg1.relation_triples_list, self.kgs.kg2.relation_triples_list, |
| self.kgs.kg1.relation_triples_set, self.kgs.kg2.relation_triples_set, |
| self.kgs.kg1.entities_list, self.kgs.kg2.entities_list, |
| self.args.batch_size, steps_task, |
| batch_queue, neighbors1, neighbors2, self.args.neg_triple_num)).start() |
| epoch_loss = 0 |
| trained_samples_num = 0 |
| for i in range(triple_steps): |
| batch_pos, batch_neg = batch_queue.get() |
| batch_loss, _ = self.session.run(fetches=[self.triple_loss, self.triple_optimizer], |
| feed_dict={self.pos_hs: [x[0] for x in batch_pos], |
| self.pos_rs: [x[1] for x in batch_pos], |
| self.pos_ts: [x[2] for x in batch_pos], |
| self.neg_hs: [x[0] for x in batch_neg], |
| self.neg_rs: [x[1] for x in batch_neg], |
| self.neg_ts: [x[2] for x in batch_neg]}) |
| trained_samples_num += len(batch_pos) |
| epoch_loss += batch_loss |
| epoch_loss /= trained_samples_num |
| random.shuffle(self.kgs.kg1.relation_triples_list) |
| random.shuffle(self.kgs.kg2.relation_triples_list) |
| print('epoch {}, avg. triple loss: {:.4f}, cost time: {:.1f}s'.format(epoch, epoch_loss, time.time() - start)) |
|
|
| def launch_align_training_1epo(self, epoch, triple_steps): |
| start = time.time() |
| epoch_loss = 0 |
| trained_samples_num = 0 |
| batch_size = 2 * len(self.kgs.train_links) // triple_steps |
| kg_training_ents2 = list(self.kgs.kg2.entities_set - set(self.kgs.train_entities2)) |
|
|
| neg_batch_size = batch_size * self.args.mapping_neg_num |
|
|
| for i in range(triple_steps): |
| links_batch = random.sample(self.kgs.train_links, batch_size) |
| pos_entities1 = random.sample(self.kgs.kg1.entities_list, neg_batch_size) |
| neg_entities2 = random.sample(kg_training_ents2, neg_batch_size) |
| batch_loss, _ = self.session.run(fetches=[self.mapping_loss, self.mapping_optimizer], |
| feed_dict={self.seed_entities1: [x[0] for x in links_batch], |
| self.seed_entities2: [x[1] for x in links_batch], |
| self.pos_entities1: pos_entities1, |
| self.neg_entities2: neg_entities2}) |
| epoch_loss += batch_loss |
| trained_samples_num += len(links_batch) |
| epoch_loss /= trained_samples_num |
| print('epoch {}, avg. mapping loss: {:.4f}, cost time: {:.1}s'.format(epoch, epoch_loss, time.time() - start)) |
|
|
| def generate_neighbors(self): |
| t1 = time.time() |
| assert 0.0 < self.args.truncated_epsilon < 1.0 |
|
|
| num1 = len(self.kgs.kg1.entities_list) // 2 |
| if len(self.kgs.kg1.entities_list) > 200000: |
| num1 = len(self.kgs.kg1.entities_list) // 3 |
| if num1 > len(self.kgs.useful_entities_list1): |
| kg1_random_ents = self.kgs.useful_entities_list1 + \ |
| random.sample(list(set(self.kgs.kg1.entities_list) - set(self.kgs.useful_entities_list1)), |
| num1 - len(self.kgs.useful_entities_list1)) |
| else: |
| kg1_random_ents = self.kgs.useful_entities_list1 |
| embeds1 = self.eval_embeddings(kg1_random_ents) |
|
|
| num2 = len(self.kgs.kg2.entities_list) // 2 |
| if len(self.kgs.kg2.entities_list) > 200000: |
| num2 = len(self.kgs.kg2.entities_list) // 3 |
| if num2 > len(self.kgs.useful_entities_list2): |
| kg2_random_ents = self.kgs.useful_entities_list2 + \ |
| random.sample(list(set(self.kgs.kg2.entities_list) - set(self.kgs.useful_entities_list2)), |
| num2 - len(self.kgs.useful_entities_list2)) |
| else: |
| kg2_random_ents = self.kgs.useful_entities_list2 |
| embeds2 = self.eval_embeddings(kg2_random_ents) |
| neighbors_num1 = int((1 - self.args.truncated_epsilon) * num1) |
| neighbors_num2 = int((1 - self.args.truncated_epsilon) * num2) |
| print("generating neighbors...") |
| neighbors1 = generate_neighbours(embeds1, kg1_random_ents, neighbors_num1, |
| frags_num=self.args.batch_threads_num) |
| neighbors2 = generate_neighbours(embeds2, kg2_random_ents, neighbors_num2, |
| frags_num=self.args.batch_threads_num) |
| print("generating neighbors ({}, {}) costs {:.3f} s.".format(num1, num2, time.time() - t1)) |
| gc.collect() |
| return neighbors1, neighbors2 |
|
|
| def run(self): |
| t = time.time() |
|
|
| triples_num = self.kgs.kg1.relation_triples_num + self.kgs.kg2.relation_triples_num |
| triple_steps = int(math.ceil(triples_num / self.args.batch_size)) |
| steps_tasks = task_divide(list(range(triple_steps)), self.args.batch_threads_num) |
| training_batch_manager = mp.Manager() |
| training_batch_queue = training_batch_manager.Queue() |
| neighbors1, neighbors2 = dict(), dict() |
|
|
| |
| for i in range(1, self.args.max_epoch + 1): |
| self.launch_training_1epo(i, triple_steps, steps_tasks, training_batch_queue, neighbors1, neighbors2) |
| |
| if i >= self.args.start_valid and i % self.args.eval_freq == 0: |
| flag = self.valid_alignment(self.args.stop_metric, False) |
| self.flag1, self.flag2, self.early_stop = early_stop(self.flag1, self.flag2, flag) |
| if self.early_stop or i == self.args.max_epoch: |
| break |
| |
| if self.args.neg_sampling == 'truncated' and i % self.args.truncated_freq == 0: |
| if neighbors1 is not None: |
| del neighbors1, neighbors2 |
| gc.collect() |
| neighbors1, neighbors2 = self.generate_neighbors() |
|
|
| print("Training ends. Total time = {:.1f} s.".format(time.time() - t)) |
|
|
