Skip to content

GitLab

Unverified Commit 3bc31098 authored by Da Zheng's avatar Da Zheng Committed by GitHub
Browse files

use PBG sampler for testing. (#899)

parent 251a9842
Hide whitespace changes
Inline Side-by-side
Showing with 8 additions and 8 deletions
apps/kg/eval.py
View file @ 3bc31098
......@@ -105,12 +105,12 @@ def main(args):
for i in range(args.num_proc):
test_sampler_head = eval_dataset.create_sampler('test', args.batch_size,
args.neg_sample_size,
mode='head',
mode='PBG-head',
num_workers=args.num_worker,
rank=i, ranks=args.num_proc)
test_sampler_tail = eval_dataset.create_sampler('test', args.batch_size,
args.neg_sample_size,
mode='tail',
mode='PBG-tail',
num_workers=args.num_worker,
rank=i, ranks=args.num_proc)
test_sampler_heads.append(test_sampler_head)
......@@ -118,12 +118,12 @@ def main(args):
else:
test_sampler_head = eval_dataset.create_sampler('test', args.batch_size,
args.neg_sample_size,
mode='head',
mode='PBG-head',
num_workers=args.num_worker,
rank=0, ranks=1)
test_sampler_tail = eval_dataset.create_sampler('test', args.batch_size,
args.neg_sample_size,
mode='tail',
mode='PBG-tail',
num_workers=args.num_worker,
rank=0, ranks=1)
......
apps/kg/train.py
View file @ 3bc31098
......@@ -208,12 +208,12 @@ def run(args, logger):
for i in range(args.num_proc):
test_sampler_head = eval_dataset.create_sampler('test', args.batch_size_eval,
args.neg_sample_size_test,
mode='head',
mode='PBG-head',
num_workers=args.num_worker,
rank=i, ranks=args.num_proc)
test_sampler_tail = eval_dataset.create_sampler('test', args.batch_size_eval,
args.neg_sample_size_test,
mode='tail',
mode='PBG-tail',
num_workers=args.num_worker,
rank=i, ranks=args.num_proc)
test_sampler_heads.append(test_sampler_head)
......@@ -221,12 +221,12 @@ def run(args, logger):
else:
test_sampler_head = eval_dataset.create_sampler('test', args.batch_size_eval,
args.neg_sample_size_test,
mode='head',
mode='PBG-head',
num_workers=args.num_worker,
rank=0, ranks=1)
test_sampler_tail = eval_dataset.create_sampler('test', args.batch_size_eval,
args.neg_sample_size_test,
mode='tail',
mode='PBG-tail',
num_workers=args.num_worker,
rank=0, ranks=1)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment