Exemplo n.º 1
0
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()
Exemplo n.º 2
0
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)
Exemplo n.º 3
0
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-----------------------------------------------')