def train(infile, outfile, seed):
print("infile:", infile)
np.random.seed(seed)
d = pickle.load(open(infile, 'rb'))
tgt_nucs = d['tgt_nucs']
MAX_N = d['MAX_N']
print("TGT_NUCS=", tgt_nucs)
print("output is", outfile)
USE_CUDA = True
mu_scale = []
std_scale = []
for tn in tgt_nucs:
for k, v in NUC_LOSS_SCALE.items():
if k in tn:
mu_scale.append(v)
for k, v in NUC_STD_SCALE.items():
if k in tn:
std_scale.append(v)
assert len(mu_scale) == len(tgt_nucs)
assert len(std_scale) == len(tgt_nucs)
print("NUC_LOSS_SCALE=", NUC_LOSS_SCALE)
print("mu_scale=", mu_scale)
print("std_scale=", std_scale)
print("tgt_nucs=", tgt_nucs)
### Create datasets and data loaders
BATCH_SIZE = 16
dataset_hparams = graph_conv_many_nuc_util.DEFAULT_DATA_HPARAMS
ds_train, ds_test = graph_conv_many_nuc_util.make_datasets(
{'filename': infile}, dataset_hparams)
dl_train = torch.utils.data.DataLoader(ds_train,
batch_size=BATCH_SIZE,
shuffle=True,
pin_memory=True)
dl_test = torch.utils.data.DataLoader(ds_test,
batch_size=BATCH_SIZE,
shuffle=True,
pin_memory=True)
net_params = net_params_base.copy()
opt_params = opt_params_base.copy()
net_params['g_feature_n'] = ds_test[0][1].shape[-1]
net_params['OUT_DIM'] = len(tgt_nucs)
use_std = True
net = nets.GraphVertModel(**net_params)
net = move(net, USE_CUDA)
for n, p in net.named_parameters():
print(n, p.shape)
loss_config = {
'std_regularize': STD_REGULARIZE,
'mu_scale': mu_scale,
'std_scale': std_scale
}
if use_std:
std_regularize = loss_config['std_regularize']
mu_scale = move(torch.Tensor(loss_config['mu_scale']), USE_CUDA)
std_scale = move(torch.Tensor(loss_config['std_scale']), USE_CUDA)
criterion = nets.NormUncertainLoss(
mu_scale, std_scale, std_regularize=loss_config['std_regularize'])
else:
criterion = nets.MaskedMSELoss()
validate_func = create_uncertain_validate_func(tgt_nucs)
optimizer = torch.optim.Adam(net.parameters(),
lr=opt_params['lr'],
amsgrad=opt_params['amsgrad'],
eps=opt_params['eps'],
weight_decay=opt_params['weight_decay'])
scheduler = torch.optim.lr_scheduler.StepLR(
optimizer,
step_size=opt_params['scheduler_step_size'],
gamma=opt_params['scheduler_gamma'])
MODEL_NAME = "{}.{:08d}".format(EXP_NAME, int(time.time() % 1e8))
checkpoint_filename = os.path.join(CHECKPOINT_DIR,
MODEL_NAME + ".{epoch_i:08d}")
print("checkpoint:", checkpoint_filename)
checkpoint_func = graph_conv_many_nuc_util.create_checkpoint_func(
10, checkpoint_filename)
writer = SummaryWriter("{}/{}".format(TENSORBOARD_DIR, MODEL_NAME))
metadata = {
'dataset_hparams': dataset_hparams,
'net_params': net_params,
'opt_params': opt_params,
'infile': infile,
'tgt_nucs': tgt_nucs,
'max_n': MAX_N,
'loss_params': loss_config
}
json.dump(metadata,
open(os.path.join(CHECKPOINT_DIR, MODEL_NAME + ".json"), 'w'),
indent=4)
print(json.dumps(metadata, indent=4))
pickle.dump(metadata,
open(os.path.join(CHECKPOINT_DIR, MODEL_NAME + ".meta"), 'wb'))
graph_conv_many_nuc_util.generic_runner(net,
optimizer,
scheduler,
criterion,
dl_train,
dl_test,
MAX_EPOCHS=MAX_EPOCHS,
USE_CUDA=USE_CUDA,
writer=writer,
validate_func=validate_func,
checkpoint_func=checkpoint_func)
pickle.dump(
{
'net_params': net_params,
'opt_params': opt_params,
'loss_params': loss_config
}, open(outfile, 'wb'))
def train(infiles, outfile, seed, cv_i):
print("infiles:", infiles)
meta_infile, data_infile = infiles
np.random.seed(seed)
config = pickle.load(open(meta_infile, 'rb'))['config']
tgt_nucs = config['spectra_nuc']
MAX_N = config['max_atom_n']
print("TGT_NUCS=", tgt_nucs)
print("output is", outfile)
USE_CUDA = True
mu_scale = []
std_scale = []
for tn in tgt_nucs:
for k, v in NUC_LOSS_SCALE.items():
if k in tn:
mu_scale.append(v)
for k, v in NUC_STD_SCALE.items():
if k in tn:
std_scale.append(v)
assert len(mu_scale) == len(tgt_nucs)
assert len(std_scale) == len(tgt_nucs)
print("NUC_LOSS_SCALE=", NUC_LOSS_SCALE)
print("mu_scale=", mu_scale)
print("std_scale=", std_scale)
print("tgt_nucs=", tgt_nucs)
### Create datasets and data loaders
BATCH_SIZE = 32
dataset_hparams = graph_conv_many_nuc_util.DEFAULT_DATA_HPARAMS.copy()
dataset_hparams['feat_vect_args']['mmff_atom_types_onehot'] = True
dataset_hparams['feat_vect_args']['feat_atomicno'] = False
# dataset_hparams['adj_args']['mat_power'] = [1,2,3,4, 6, 8]
# dataset_hparams['adj_args']['split_weights'] = None
# dataset_hparams['adj_args']['edge_weighted'] = False
ds_train, ds_test = graph_conv_many_nuc_util.make_datasets({'filename' : data_infile},
dataset_hparams, MAX_N, cv_i=cv_i)
dl_train = torch.utils.data.DataLoader(ds_train, batch_size=BATCH_SIZE,
shuffle=True,pin_memory=True)
dl_test = torch.utils.data.DataLoader(ds_test, batch_size=BATCH_SIZE,
shuffle=True,pin_memory=True)
net_params = net_params_base.copy()
opt_params = opt_params_base.copy()
net_params['vert_f_in'] = ds_test[0][1].shape[-1]
net_params['OUT_DIM'] = len(tgt_nucs)
net_params['MAX_N'] = MAX_N
#net_name = 'nets.SGCModel'
net_name = 'nets.RelNetFromS2S'
net = eval(net_name)(**net_params)
net = move(net, USE_CUDA)
for n, p in net.named_parameters():
print(n, p.shape)
if USE_STD:
loss_config = {'std_regularize' : STD_REGULARIZE,
'mu_scale' : mu_scale,
'std_scale' : std_scale}
std_regularize = loss_config['std_regularize']
mu_scale = move(torch.Tensor(loss_config['mu_scale']), USE_CUDA)
std_scale = move(torch.Tensor(loss_config['std_scale']), USE_CUDA)
criterion = nets.NormUncertainLoss(mu_scale,
std_scale,
std_regularize = loss_config['std_regularize'])
else:
loss_config = {'norm' : 'huber',
'scale' : 0.1}
criterion = nets.NoUncertainLoss(**loss_config)
opt_direct_params = {}
optimizer_name = opt_params.get('optimizer', 'adam')
if optimizer_name == 'adam':
for p in ['lr', 'amsgrad', 'eps', 'weight_decay', 'momentum']:
if p in opt_params:
opt_direct_params[p] = opt_params[p]
optimizer = torch.optim.Adam(net.parameters(), **opt_direct_params)
elif optimizer_name == 'adamax':
for p in ['lr', 'eps', 'weight_decay', 'momentum']:
if p in opt_params:
opt_direct_params[p] = opt_params[p]
optimizer = torch.optim.Adamax(net.parameters(), **opt_direct_params)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=opt_params['scheduler_step_size'],
gamma=opt_params['scheduler_gamma'])
MODEL_NAME = "{}.{:08d}".format(EXP_NAME,int(time.time() % 1e8))
checkpoint_filename = os.path.join(CHECKPOINT_DIR, MODEL_NAME + ".{epoch_i:08d}")
print("checkpoint:", checkpoint_filename)
checkpoint_func = graph_conv_many_nuc_util.create_checkpoint_func(CHECKPOINT_EVERY, checkpoint_filename)
writer = SummaryWriter("{}/{}".format(TENSORBOARD_DIR, MODEL_NAME))
validate_func = create_uncertain_validate_func(tgt_nucs, writer)
metadata = {'dataset_hparams' : dataset_hparams,
'net_params' : net_params,
'opt_params' : opt_params,
'meta_infile' : meta_infile,
'data_infile' : data_infile,
'tgt_nucs' : tgt_nucs,
'max_n' : MAX_N,
'net_name' : net_name,
'batch_size' : BATCH_SIZE,
'loss_params' : loss_config}
json.dump(metadata, open(os.path.join(CHECKPOINT_DIR, MODEL_NAME + ".json"), 'w'),
indent=4)
print(json.dumps(metadata, indent=4))
print("MODEL_NAME=", MODEL_NAME)
pickle.dump(metadata,
open(os.path.join(CHECKPOINT_DIR, MODEL_NAME + ".meta"), 'wb'))
graph_conv_many_nuc_util.generic_runner(net, optimizer, scheduler, criterion,
dl_train, dl_test,
MAX_EPOCHS=MAX_EPOCHS,
USE_CUDA=USE_CUDA, writer=writer,
validate_func= validate_func,
checkpoint_func= checkpoint_func)
pickle.dump({'net_params' : net_params,
'opt_params' : opt_params,
'loss_params' : loss_config},
open(outfile, 'wb'))