def run_grid(nentity,
nrelation,
train_triples,
valid_triples,
test_triples,
all_true_triples,
args,
rule_iterators=None,
adv_model=None):
ntriples = len(train_triples)
if args.inject:
print('injecting rules')
else:
print('rules not injected')
if args.ruge:
print('Using RUGE injection model')
reset_empty_values(args)
current_learning_rate = args.learning_rate
if args.negative_adversarial_sampling:
print('Temperature - ', args.adversarial_temperature)
print()
info = 'Model - {}; opt - {}; batch size - {}; dataset - {}; lr - {}, gamma = {}; '.format(
args.model, args.opt, args.batch_size, args.data_path,
current_learning_rate, args.gamma)
info2 = 'Loss fnc - {}; inv - {}; impl - {}; sym - {}; eq - {}'.format(
args.loss, args.inv, args.impl, args.sym, args.eq)
print(info)
print(info2)
current_learning_rate = args.learning_rate
EPSILONS = itertools.product(EPSILONS_INV, EPSILONS_IMPL, EPSILONS_SYM,
EPSILONS_EQ)
WEIGHTS = itertools.product(WEIGHTS_INV, WEIGHTS_IMPL, WEIGHTS_SYM,
WEIGHTS_EQ)
idx = -1 # for saving models with several parameters
for g1, g2 in zip(GAMMA1, GAMMA2):
for eps_inv, eps_impl, eps_sym, eps_eq in EPSILONS:
for w_inv, w_impl, w_sym, w_eq in WEIGHTS:
for dim, n_negs, steps in itertools.product(
DIMENSIONS, N_NEGS_LIST, N_STEPS_LIST):
idx += 1
# re-initialize the model
kge_model = KGEModel(model_name=args.model,
nentity=nentity,
nrelation=nrelation,
ntriples=ntriples,
hidden_dim=dim,
args=args)
if 'inverse' in RULE_TYPES:
kge_model.rule_weight['inverse'] = w_inv
kge_model.epsilon_inv = eps_inv
if 'implication' in RULE_TYPES:
kge_model.rule_weight['implication'] = w_impl
kge_model.epsilon_impl = eps_impl
if 'symmetry' in RULE_TYPES:
kge_model.rule_weight['symmetry'] = w_sym
kge_model.epsilon_sym = eps_sym
if 'equality' in RULE_TYPES:
kge_model.rule_weight['equality'] = w_eq
kge_model.epsilon_eq = eps_eq
kge_model.set_loss(args.loss)
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('optimizer: %s' % args.opt)
logging.info('learning rate: %f' % current_learning_rate)
logging.info('loss: %s' % args.loss)
if args.inv:
logging.info(
'using inverse rules: eps = %f, weight = %f' %
(kge_model.epsilon_inv,
kge_model.rule_weight['inverse']))
if args.impl:
logging.info(
'using implication rules: eps = %f, weight = %f' %
(kge_model.epsilon_impl,
kge_model.rule_weight['implication']))
if args.sym:
logging.info(
'using symmetry rules: eps = %f, weight = %f' %
(kge_model.epsilon_sym,
kge_model.rule_weight['symmetry']))
if args.eq:
logging.info(
'using equality rules: eps = %f, weight = %f' %
(kge_model.epsilon_eq,
kge_model.rule_weight['equality']))
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:
kge_model = kge_model.cuda()
logging.info('Ramdomly Initializing %s Model...' %
args.model)
print_rules_info(kge_model, args)
args.max_steps = steps
args.negative_sample_size = n_negs
#out_line = '#steps = {}, #negs = {};'.format(args.max_steps, args.negative_sample_size)
logging.info('Max steps - %d' % args.max_steps)
logging.info('Negative sample %d ' %
args.negative_sample_size)
assert kge_model.inject == args.inject, 'Inject is wrong'
# train
train_iterator = construct_dataloader(
args, train_triples, nentity, nrelation)
step = train_model(0, valid_triples, all_true_triples,
kge_model, adv_model, train_iterator,
rule_iterators, args, str(idx))
# valid
logging.info('Evaluating on Valid Dataset...')
metrics = kge_model.test_step(kge_model, valid_triples,
all_true_triples, args)
#metrics1 = kge_model.getScore(kge_model, valid_triples, all_true_triples, args)
log_metrics('Valid', step, metrics)
info = 'Validation (%d): ' % step
for key, val in metrics.items():
info = info + key + ' - ' + str(val) + ';'
print(info)
# test
out_line = '#steps = {}, #negs = {}, dim = {};'.format(
step, args.negative_sample_size, kge_model.hidden_dim)
metrics = kge_model.test_step(kge_model, test_triples,
all_true_triples, args)
print("Hello")
metrics1 = kge_model.getScore(kge_model, test_triples,
all_true_triples, args)
log_metrics('Test', step, metrics)
values = [
str(metrics['MRR']),
str(metrics['MR']),
str(metrics['HITS@1']),
str(metrics['HITS@3']),
str(metrics['HITS@10'])
]
out_line = out_line + ';'.join(values)
print(out_line)
logging.info(
'\n-----------------------------------------------')
print()
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 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(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
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
current_learning_rate = args.learning_rate
ntriples = len(train_triples)
'''print('Model: %s' % args.model)
print('Data Path: %s' % args.data_path)
print('#entity: %d' % nentity)
print('#relation: %d' % nrelation)
print('optimizer: ', OPT)
if args.train_old: print('USING ORIGINAL TRAINING FUNCTION')
print('learning_rate = %f' % current_learning_rate)
print('batch_size = %d' % args.batch_size)
print('negative_adversarial_sampling = %d' % args.negative_adversarial_sampling)
print('hidden_dim = %d' % args.hidden_dim)
print('negative_adversarial_sampling = %s' % str(args.negative_adversarial_sampling))
if args.negative_adversarial_sampling:
print('adversarial_temperature = %f' % args.adversarial_temperature)
'''
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('optimizer: %s' % OPT)
if args.train_old: logging.info('USING ORIGINAL TRAINING FUNCTION')
#else: print('GRID TESTING\nUsing new loss function')
info = 'Model - {}; opt - {}; batch size - {}, dim - {}; dataset - {}; lr - {}; '.format(
args.model, OPT, str(args.batch_size), args.hidden_dim, args.data_path,
str(current_learning_rate))
print(info)
for g1, g2 in zip(GAMMA1, GAMMA2):
for n_neg in N_NEGS_LIST:
for steps in N_STEPS_LIST:
current_learning_rate = args.learning_rate
# re-initialize the model
kge_model = KGEModel(
model_name=args.model,
nentity=nentity,
nrelation=nrelation,
ntriples=ntriples,
hidden_dim=args.hidden_dim,
gamma=args.gamma,
gamma1=g1,
gamma2=g2,
double_entity_embedding=args.double_entity_embedding,
double_relation_embedding=args.double_relation_embedding,
)
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:
kge_model = kge_model.cuda()
logging.info('Ramdomly Initializing %s Model...' % args.model)
args.max_steps = steps
args.negative_sample_size = n_neg
out_line = 'g1 = {}, g2 = {}, #steps = {}, #negs = {};'.format(
kge_model.gamma1, kge_model.gamma2, args.max_steps,
args.negative_sample_size)
logging.info('gamma1 = %f, gamma2 = %f' % (g1, g2))
logging.info('Max steps - %d' % args.max_steps)
logging.info('Negative sample %d ' % args.negative_sample_size)
train_iterator = construct_dataloader(args, train_triples,
nentity, nrelation)
step = grid_train_model(0, valid_triples, all_true_triples,
kge_model, train_iterator, args)
metrics = kge_model.test_step(kge_model, test_triples,
all_true_triples, args)
log_metrics('Test', step, metrics)
values = [
str(metrics['MRR']),
str(metrics['MR']),
str(metrics['HITS@1']),
str(metrics['HITS@3']),
str(metrics['HITS@10'])
]
out_line = out_line + ';'.join(values)
print(out_line)
logging.info(
'\n-----------------------------------------------')