def main(args):
if args.gpu < 0:
cuda = False
else:
cuda = True
torch.cuda.set_device(args.gpu)
default_path = create_default_path()
print('\n*** Set default saving/loading path to:', default_path)
if args.dataset == AIFB or args.dataset == MUTAG:
module = importlib.import_module(MODULE.format('dglrgcn'))
data = module.load_dglrgcn(args.data_path)
data = to_cuda(data) if cuda else data
mode = NODE_CLASSIFICATION
elif args.dataset == MUTAGENICITY or args.dataset == PTC_MR or args.dataset == PTC_MM or args.dataset == PTC_FR or args.dataset == PTC_FM:
module = importlib.import_module(MODULE.format('dortmund'))
data = module.load_dortmund(args.data_path)
data = to_cuda(data) if cuda else data
mode = GRAPH_CLASSIFICATION
else:
raise ValueError('Unable to load dataset', args.dataset)
print_graph_stats(data[GRAPH])
config_params = read_params(args.config_fpath, verbose=True)
# create GNN model
model = Model(g=data[GRAPH],
config_params=config_params[0],
n_classes=data[N_CLASSES],
n_rels=data[N_RELS] if N_RELS in data else None,
n_entities=data[N_ENTITIES] if N_ENTITIES in data else None,
is_cuda=cuda,
mode=mode)
if cuda:
model.cuda()
# 1. Training
app = App()
learning_config = {
'lr': args.lr,
'n_epochs': args.n_epochs,
'weight_decay': args.weight_decay,
'batch_size': args.batch_size,
'cuda': cuda
}
print('\n*** Start training ***\n')
app.train(data, config_params[0], learning_config, default_path, mode=mode)
# 2. Testing
print('\n*** Start testing ***\n')
app.test(data, default_path, mode=mode)
# 3. Delete model
remove_model(default_path)
def prepare_files(self, task_ids=None, cuda=True):
data = self.prep_data.load_data_files(task_ids)
if data is None:
return None, self.args
data = to_cuda(data) if cuda is True else data
return data, self.args
def load_dataset(args, cuda):
if 'prep_path' in config_params and config_params['prep_path'] is not None:
sys.path.insert(0, config_params['prep_path'])
print('*** prep path', config_params['prep_path'])
from prep_data_n import PrepareData
else:
from utils.prep_data_n import PrepareData
prep_data = PrepareData(args, cuckoo_analysis_dir='api_tasks/data_report')
data = prep_data.load_data()
data = to_cuda(data) if cuda is True else data
# if cuda is True:
# data[GRAPH] = data[GRAPH].to(torch.device('cuda:0'))
return data
def load_dataset(args, cuda):
prep_data = PrepareData(reports_parent_dir_path=args.input_report_folder,
data_json_path=args.input_data_file,
pickle_folder=args.input_data_folder,
vocab_path=args.vocab_path,
encode_edge_data=args.encode_edge_data,
save_json=args.save_json,
use_interesting_apis=args.use_interesting_apis,
prepend_vocab=args.prepend_vocab,
mapping_path=args.mapping_path)
data = prep_data.load_data(from_folder=args.from_report_folder,
from_json=args.from_data_json,
from_pickle=args.from_pickle)
data = to_cuda(data) if cuda else data
return data
def prepare_files(self, cuda=True):
if self.task_ids is None:
data = self.prep_data.load_data_files(
self.task_ids,
report_dir_path=self.report_dir_path,
report_dir_name=self.report_dir_name,
report_file_name=self.report_file_name)
else:
data = self.prep_data.load_data_files(self.task_ids)
# data = self.prep_data.load_data()
if data is None:
return None, self.args
data = to_cuda(data) if cuda is True else data
return data
def prepare_files(task_ids=None, cuda=True):
args = read_params(CONFIG_PATH, verbose=False)
args['config_fpath'] = CONFIG_PATH
if task_ids is not None:
task_ids = [str(tid) for tid in task_ids]
batch_task_name = '-'.join(task_ids)
args["input_report_folder"] = cf.__ROOT__ + '/api_tasks/data_report'
args[
"input_data_folder"] = cf.__ROOT__ + '/api_tasks/data_json/{}'.format(
batch_task_name)
args[
"input_pickle_folder"] = cf.__ROOT__ + '/api_tasks/data_pickle/{}'.format(
batch_task_name)
args["mapping_path"] = cf.__ROOT__ + '/' + args["mapping_path"]
args["train_embedding_path"] = cf.__ROOT__ + '/' + args[
"train_embedding_path"]
args["vocab_path"] = cf.__ROOT__ + '/' + args["vocab_path"]
args["graph_viz_dir"] = cf.__ROOT__ + '/data/graphviz'
args["from_pickle"] = False
args["from_report_folder"] = False
args["from_data_json"] = False
args["prepare_word_embedding"] = True
args["train_embedder"] = False
args["reverse_edge"] = __REVERSE_EDGE__
del args['train_list_file']
del args['test_list_file']
args['do_draw'] = __DO_DRAW__
prep_data = PrepareData(args)
data = prep_data.load_data_files(task_ids)
if data is None:
return None, args
data = to_cuda(data) if cuda is True else data
return data, args