logging.basicConfig(level=logging.DEBUG,
format='%(module)5s %(asctime)s %(message)s',
datefmt='%H:%M:%S')
parser = argparse.ArgumentParser(description='GAN for Knowledge Embedding')
parser.add_argument('--config',
type=str,
default="pretrain_wn18_50",
metavar='S',
help='config file')
param = parser.parse_args()
args = Config(param.config)
args.cuda = not args.no_cuda and torch.cuda.is_available()
args.device = torch.device("cuda" if args.cuda else "cpu")
if args.cuda:
torch.cuda.set_device(0)
logging.getLogger().addHandler(
logging.FileHandler("log/" + args.perfix +
datetime.datetime.now().strftime("%m%d%H%M%S")))
# init dataset
train_data = KBDataset(os.path.join("data", args.task_dir, 'train2id.txt'),
args.cuda)
valid_data = KBDataset(os.path.join("data", args.task_dir, 'valid2id.txt'),
args.cuda)
test_data = KBDataset(os.path.join("data", args.task_dir, 'test2id.txt'),
args.cuda)
chk = DataChecker(args.task_dir)
