batch_size=parameters['batch_size'],
shuffle=True)
val_pos_loader = data.DataLoader(torch.Tensor(
np.array(kg.get_data_by_type('val'))),
batch_size=parameters['batch_size'],
shuffle=False)
test_pos_loader = data.DataLoader(torch.Tensor(
np.array(kg.get_data_by_type('test'))),
batch_size=parameters['batch_size'],
shuffle=False)
# Neg samplers
train_neg_sampler = RandomNegativeSampler(
kg.get_num_of_ent("train"),
kg.get_num_of_ent('train'),
neg_per_pos=parameters['num_of_neg_samples'])
if args.random_val_neg_sampler:
val_neg_sampler = RandomNegativeSampler(
kg.get_num_of_ent("val"),
kg.get_num_of_ent('val'),
neg_per_pos=parameters['num_of_neg_samples'])
else:
val_neg_sampler = HonestNegativeSampler(kg.get_num_of_ent("val"),
kg.get_num_of_ent('val'),
kg.df_drug,
neg_per_pos=1)
test_neg_sampler = HonestNegativeSampler(kg.get_num_of_ent("test"),
parameters['regularization'] = args.regularization
parameters['use_proteins'] = args.use_proteins
parameters['reversed'] = args.reversed
parameters['metrics_separately'] = args.metrics_separately
parameters['random_val_neg_sampler'] = args.random_val_neg_sampler
parameters['val_regenerate'] = args.val_regenerate
use_cuda = args.gpu and torch.cuda.is_available()
device = torch.device("cuda" if args.gpu else "cpu")
print(f'Use device: {device}')
kg = KnowledgeGraph(data_path=DATA_CONST['work_dir'],
use_proteins=args.use_proteins,
use_proteins_on_validation=False,
use_reversed_edges=args.reversed)
model = TriVec(ent_total=kg.get_num_of_ent('train'),
rel_total=kg.get_num_of_rel('train'))
model = model.to(device)
loss_func = NegativeSoftPlusLoss()
print('Test')
checkpoint = torch.load(args.model_path)
model.load_state_dict(checkpoint['model_state_dict'])
switch_grad_mode(model, requires_grad=False)
switch_model_mode(model, train=False)
model.eval()
test_trivec_model('seen')
print("===========================")
test_trivec_model('new')