'reduction': 4
},
exp_dir=fold_dir,
chk_score_attr='rank_score',
log_params={
'n_gen': '#gen',
'C': 'C',
'depth': 'Depth'
},
gpus=gpus,
gpu_threshold=0.75)
if fold_idx in args.test:
best_dict = get_best_info(os.path.join(fold_dir, 'design'),
mode='manual')
t_config = best_dict['config']
ts_ld = DataLoader(dataset[ts_i.tolist()],
batch_size=512,
shuffle=False)
best_dict['ts_loss'], best_dict['ts_score'] = run_test(
trial_dir=best_dict['trial_dir'],
ts_ld=ts_ld,
model_func=lambda config: CGMN(config['out'], config[
'n_gen'], config['C'], config['symbols'], config['depth']),
loss_fn=get_loss_fn('bce'),
score_fn=get_score_fn('accuracy', t_config['out']),
gpus=[])
print(best_dict)
torch.save(best_dict, os.path.join(fold_dir, 'test_res.pkl'))
'batch_size': 'Batch'
},
gpus=gpus,
gpu_threshold=0.9)
if fold_idx in args.test:
best_dict = get_best_info(os.path.join(exp_dir,
f'fold_{fold_idx}'),
mode='manual')
t_config = best_dict['config']
ts_ld = DataLoader(
TreeDataset(data=[dataset[curr] for curr in ts_i]),
collate_fn=trees_collate_fn,
batch_size=512,
shuffle=False)
best_dict['ts_loss'], best_dict['ts_score'] = run_test(
trial_dir=best_dict['trial_dir'],
ts_ld=ts_ld,
model_func=lambda config: HTMN(
config['out'], math.ceil(config['n_gen'] / 2),
math.floor(config['n_gen'] / 2), config['C'], config[
'L'], config['M']),
loss_fn=get_loss_fn('bce', None),
score_fn=get_score_fn('roc-auc', t_config['out']),
gpus=[])
print(best_dict)
torch.save(
best_dict,
os.path.join(exp_dir, f'fold_{fold_idx}', 'test_res.pkl'))
def _wrapper_init_fn(config):
if config['model'] == 'htmn':
from data.tree.utils import TreeDataset, trees_collate_fn
from torch.utils.data import DataLoader
from htmn.htmn import HTMN
dataset = TreeDataset(config['wdir'], config['dataset'])
tr_ld = DataLoader(
TreeDataset(data=[dataset[i] for i in config['tr_idx']]),
batch_size=config['batch_size'],
shuffle=True,
collate_fn=trees_collate_fn,
drop_last=len(config['tr_idx']) % config['batch_size'] == 1)
if config['vl_idx'] is not None:
vl_ld = DataLoader(
TreeDataset(data=[dataset[i] for i in config['vl_idx']]),
batch_size=config['batch_size'],
shuffle=False,
collate_fn=trees_collate_fn,
drop_last=len(config['vl_idx']) % config['batch_size'] == 1)
else:
vl_ld = None
model = HTMN(config['out'], config['mode'], config['n_gen'],
config['C'], config['L'], config['M'])
opt = torch.optim.Adam(model.parameters(), lr=config['lr'])
if config['model'] == 'ghtmn':
from data.graph.g2t import ParallelTUDataset, TreeCollater, pre_transform, transform
from torch.utils.data import DataLoader
from graph_htmn.graph_htmn import GraphHTMN
tr_idx, vl_idx = config['tr_idx'], config['vl_idx']
dataset = ParallelTUDataset(
os.path.join(config['wdir'], config['dataset'],
f'D{config["depth"]}'),
config['dataset'],
pre_transform=pre_transform(config['depth']),
transform=transform(config['dataset']))
dataset.data.x = dataset.data.x.argmax(1).detach()
tr_ld = DataLoader(dataset[tr_idx],
collate_fn=TreeCollater(config['depth']),
batch_size=config['batch_size'],
shuffle=True)
vl_ld = DataLoader(dataset[vl_idx],
collate_fn=TreeCollater(config['depth']),
batch_size=config['batch_size'],
shuffle=False)
if config['gen_mode'] == 'bu':
n_bu, n_td = config['n_gen'], 0
elif config['gen_mode'] == 'td':
n_bu, n_td = 0, config['n_gen']
elif config['gen_mode'] == 'both':
n_bu, n_td = math.ceil(config['n_gen'] / 2), math.floor(
config['n_gen'] / 2)
model = GraphHTMN(config['out'], n_bu, n_td, config['C'],
config['symbols'])
opt = torch.optim.Adam(model.parameters(), lr=config['lr'])
if config['model'] == 'cgmn':
from torch_geometric.datasets import TUDataset
from torch_geometric.data import DataLoader
from cgmn.cgmn import CGMN
def tf_func(data):
data.y = data.y.unsqueeze(1)
return data
dataset = TUDataset(config['wdir'],
config['dataset'],
transform=tf_func)
dataset.data.x = dataset.data.x.argmax(1).detach()
tr_ld = DataLoader(dataset[config['tr_idx']],
batch_size=config['batch_size'],
shuffle=True,
drop_last=len(config['tr_idx']) %
config['batch_size'] == 1)
if config['vl_idx'] is not None:
vl_ld = DataLoader(dataset[config['vl_idx']],
batch_size=config['batch_size'],
shuffle=False,
drop_last=len(config['vl_idx']) %
config['batch_size'] == 1)
else:
vl_ld = None
model = CGMN(config['out'], config['n_gen'], config['C'],
config['symbols'], config['depth'])
opt = torch.optim.Adam(model.parameters(), lr=config['lr'])
loss_fn = get_loss_fn(config['loss'])
score_fn = get_score_fn(config['score'], config['out'])
rank_fn = get_rank_fn(config['rank'])
return model, opt, tr_ld, vl_ld, loss_fn, score_fn, rank_fn
if args.test == 'retrain':
best_dict = get_best_info(os.path.join(exp_dir, 'design'))
t_config = best_dict['config']
ts_loss = []
ts_acc = []
test_dir = os.path.join(exp_dir, 'test')
for i, t_dir in enumerate(os.listdir(test_dir)):
trial_dir = os.path.join(test_dir, t_dir)
if os.path.isdir(trial_dir):
v_loss, v_acc = run_test(
trial_dir=trial_dir,
ts_ld=ts_ld,
model_func=lambda config: HTMN(config['out'], config[
'mode'], config['n_gen'], config['C'], config['L'],
config['M']),
loss_fn=get_loss_fn('ce'),
score_fn=get_score_fn('accuracy', t_config['out']),
gpus=gpus)
ts_loss.append(v_loss)
ts_acc.append(v_acc)
best_dict = get_best_info(os.path.join(exp_dir, 'design'))
best_dict['ts_loss'], best_dict['ts_score'] = ts_loss, ts_acc
del best_dict['config']['tr_idx'], best_dict['config']['vl_idx']
print(best_dict)
torch.save(best_dict, os.path.join(exp_dir, 'test_res.pkl'))
if args.test == 'design':
best_dict = get_best_info(os.path.join(exp_dir, 'design'))
t_config = best_dict['config']
ts_loss, ts_acc = run_test(