def main(input_args=None):
# Parse the arguments.
args = parse_arguments(input_args)
device = args.gpu
method = args.method
if args.data_name == 'suzuki':
datafile = 'data/suzuki_type_test_v2.csv'
class_num = 119
class_dict = {'M': 28, 'L': 23, 'B': 35, 'S': 10, 'A': 17}
dataset_filename = 'test_data.npz'
labels = ['Yield', 'M', 'L', 'B', 'S', 'A', 'id']
elif args.data_name == 'CN':
datafile = 'data/CN_coupling_test.csv'
class_num = 206
class_dict = {'M': 44, 'L': 47, 'B': 13, 'S': 22, 'A': 74}
dataset_filename = 'test_CN_data.npz'
labels = ['Yield', 'M', 'L', 'B', 'S', 'A', 'id']
elif args.data_name == 'Negishi':
datafile = 'data/Negishi_test.csv'
class_num = 106
class_dict = {'M': 32, 'L': 20, 'T': 8, 'S': 10, 'A': 30}
dataset_filename = 'test_Negishi_data.npz'
labels = ['Yield', 'M', 'L', 'T', 'S', 'A', 'id']
elif args.data_name == 'PKR':
datafile = 'data/PKR_test.csv'
class_num = 83
class_dict = {
'M': 18,
'L': 6,
'T': 7,
'S': 15,
'A': 11,
'G': 1,
'O': 13,
'P': 4,
'other': 1
}
dataset_filename = 'test_PKR_data.npz'
labels = [
'Yield', 'M', 'L', 'T', 'S', 'A', 'G', 'O', 'P', 'other', 'id'
]
else:
raise ValueError('Unexpected dataset name')
cache_dir = os.path.join('input', '{}_all'.format(method))
# Dataset preparation.
def postprocess_label(label_list):
return numpy.asarray(label_list, dtype=numpy.float32)
print('Preprocessing dataset...')
# Load the cached dataset.
dataset_cache_path = os.path.join(cache_dir, dataset_filename)
dataset = None
if os.path.exists(dataset_cache_path):
print('Loading cached dataset from {}.'.format(dataset_cache_path))
dataset = NumpyTupleDataset.load(dataset_cache_path)
if dataset is None:
if args.method == 'mpnn':
preprocessor = preprocess_method_dict['ggnn']()
else:
preprocessor = preprocess_method_dict[args.method]()
parser = CSVFileParser(
preprocessor,
postprocess_label=postprocess_label,
labels=labels,
smiles_col=['Reactant1', 'Reactant2', 'Product'],
label_dicts=class_dict)
dataset = parser.parse(datafile)['dataset']
# Cache the laded dataset.
if not os.path.exists(cache_dir):
os.makedirs(cache_dir)
NumpyTupleDataset.save(dataset_cache_path, dataset)
labels = dataset.get_datasets()[-2]
ids = dataset.get_datasets()[-1][:, 1].reshape(-1, 1)
yields = dataset.get_datasets()[-1][:, 0].reshape(-1, 1).astype(
'float32') # [:,0] added
dataset = NumpyTupleDataset(*(dataset.get_datasets()[:-2] + (
yields,
labels,
)))
# Load the standard scaler parameters, if necessary.
scaler = None
test = dataset
print('Predicting...')
# Set up the regressor.
model_path = os.path.join(args.in_dir, args.model_filename)
if os.path.exists(model_path):
classifier = Classifier.load_pickle(model_path, device=args.gpu)
else:
predictor = set_up_predictor(args.method, args.unit_num,
args.conv_layers, class_num)
classifier = Classifier(predictor,
lossfun=F.sigmoid_cross_entropy,
metrics_fun=F.binary_accuracy,
device=args.gpu)
if args.load_modelname:
serializers.load_npz(args.load_modelname, classifier)
scaled_predictor = ScaledGraphConvPredictor(
graph_conv=classifier.predictor.graph_conv,
mlp=classifier.predictor.mlp)
classifier.predictor = scaled_predictor
# This callback function extracts only the inputs and discards the labels.
def extract_inputs(batch, device=None):
return concat_mols(batch, device=device)[:-1]
# Predict the output labels.
# Prediction function rewrite!!!
y_pred = classifier.predict(test, converter=extract_inputs)
y_pred_max = numpy.argmax(y_pred, axis=1)
y_pred_max = y_pred_max.reshape(-1, 1)
# y_pred_idx = y_pred.argsort(axis=1) # ascending
# Extract the ground-truth labels.
t = concat_mols(test, device=-1)[-1] # device 11/14 memory issue
original_t = cuda.to_cpu(t)
t_idx = original_t.squeeze(1)
t_idx = t_idx.argsort(axis=1)
# gt_indx = numpy.where(original_t == 1)
# Construct dataframe.
df_dict = {}
for i, l in enumerate(labels[:1]):
df_dict.update({
'y_pred_{}'.format(l): y_pred_max[:, -1].tolist(), # [:,-1]
't_{}'.format(l): t_idx[:, -1].tolist(),
})
df = pandas.DataFrame(df_dict)
# Show a prediction/ground truth table with 5 random examples.
print(df.sample(5))
n_eval = 10
for target_label in range(y_pred_max.shape[1]):
label_name = labels[:1][0][target_label]
print('label_name = {}, y_pred = {}, t = {}'.format(
label_name, y_pred_max[:n_eval, target_label], t_idx[:n_eval, -1]))
# Perform the prediction.
print('Evaluating...')
test_iterator = SerialIterator(test, 16, repeat=False, shuffle=False)
eval_result = Evaluator(test_iterator,
classifier,
converter=concat_mols,
device=args.gpu)()
print('Evaluation result: ', eval_result)
with open(os.path.join(args.in_dir, 'eval_result.json'), 'w') as f:
json.dump(eval_result, f)
res_dic = {}
for i in range(len(y_pred)):
res_dic[i] = str(ids[i])
json.dump(res_dic, open(os.path.join(args.in_dir, "test_ids.json"), "w"))
pickle.dump(y_pred, open(os.path.join(args.in_dir, "pred.pkl"), "wb"))
pickle.dump(original_t, open(os.path.join(args.in_dir, "gt.pkl"), "wb"))
def main():
# Parse the arguments.
args = parse_arguments()
# Set up some useful variables that will be used later on.
method = args.method
if args.label != 'all':
label = args.label
cache_dir = os.path.join('input', '{}_{}'.format(method, label))
labels = [label]
else:
labels = D.get_qm9_label_names()
cache_dir = os.path.join('input', '{}_all'.format(method))
# Get the filename corresponding to the cached dataset, based on the amount
# of data samples that need to be parsed from the original dataset.
num_data = args.num_data
if num_data >= 0:
dataset_filename = 'data_{}.npz'.format(num_data)
else:
dataset_filename = 'data.npz'
# Load the cached dataset.
dataset_cache_path = os.path.join(cache_dir, dataset_filename)
dataset = None
if os.path.exists(dataset_cache_path):
print('Loading cached data from {}.'.format(dataset_cache_path))
dataset = NumpyTupleDataset.load(dataset_cache_path)
if dataset is None:
print('Preprocessing dataset...')
preprocessor = preprocess_method_dict[method]()
dataset = D.get_qm9(preprocessor, labels=labels)
# Cache the newly preprocessed dataset.
if not os.path.exists(cache_dir):
os.mkdir(cache_dir)
NumpyTupleDataset.save(dataset_cache_path, dataset)
# Use a predictor with scaled output labels.
model_path = os.path.join(args.in_dir, args.model_filename)
regressor = Regressor.load_pickle(model_path, device=args.gpu)
scaler = regressor.predictor.scaler
if scaler is not None:
scaled_t = scaler.transform(dataset.get_datasets()[-1])
dataset = NumpyTupleDataset(*(dataset.get_datasets()[:-1] +
(scaled_t, )))
# Split the dataset into training and testing.
train_data_size = int(len(dataset) * args.train_data_ratio)
_, test = split_dataset_random(dataset, train_data_size, args.seed)
# This callback function extracts only the inputs and discards the labels.
def extract_inputs(batch, device=None):
return concat_mols(batch, device=device)[:-1]
def postprocess_fn(x):
if scaler is not None:
scaled_x = scaler.inverse_transform(x)
return scaled_x
else:
return x
# Predict the output labels.
print('Predicting...')
y_pred = regressor.predict(test,
converter=extract_inputs,
postprocess_fn=postprocess_fn)
# Extract the ground-truth labels.
t = concat_mols(test, device=-1)[-1]
original_t = scaler.inverse_transform(t)
# Construct dataframe.
df_dict = {}
for i, l in enumerate(labels):
df_dict.update({
'y_pred_{}'.format(l): y_pred[:, i],
't_{}'.format(l): original_t[:, i],
})
df = pandas.DataFrame(df_dict)
# Show a prediction/ground truth table with 5 random examples.
print(df.sample(5))
n_eval = 10
for target_label in range(y_pred.shape[1]):
label_name = labels[target_label]
diff = y_pred[:n_eval, target_label] - original_t[:n_eval,
target_label]
print('label_name = {}, y_pred = {}, t = {}, diff = {}'.format(
label_name, y_pred[:n_eval, target_label],
original_t[:n_eval, target_label], diff))
# Run an evaluator on the test dataset.
print('Evaluating...')
test_iterator = SerialIterator(test, 16, repeat=False, shuffle=False)
eval_result = Evaluator(test_iterator,
regressor,
converter=concat_mols,
device=args.gpu)()
print('Evaluation result: ', eval_result)
# Save the evaluation results.
save_json(os.path.join(args.in_dir, 'eval_result.json'), eval_result)
# Calculate mean abs error for each label
mae = numpy.mean(numpy.abs(y_pred - original_t), axis=0)
eval_result = {}
for i, l in enumerate(labels):
eval_result.update({l: mae[i]})
save_json(os.path.join(args.in_dir, 'eval_result_mae.json'), eval_result)
def main():
# Parse the arguments.
args = parse_arguments()
if args.label:
labels = args.label
class_num = len(labels) if isinstance(labels, list) else 1
else:
raise ValueError('No target label was specified.')
# Dataset preparation. Postprocessing is required for the regression task.
def postprocess_label(label_list):
return numpy.asarray(label_list, dtype=numpy.float32)
# Apply a preprocessor to the dataset.
print('Preprocessing dataset...')
preprocessor = preprocess_method_dict[args.method]()
parser = CSVFileParser(preprocessor,
postprocess_label=postprocess_label,
labels=labels,
smiles_col='SMILES')
dataset = parser.parse(args.datafile)['dataset']
# Scale the label values, if necessary.
if args.scale == 'standardize':
scaler = StandardScaler()
labels = scaler.fit_transform(dataset.get_datasets()[-1])
dataset = NumpyTupleDataset(*(dataset.get_datasets()[:-1] +
(labels, )))
else:
scaler = None
# Split the dataset into training and validation.
train_data_size = int(len(dataset) * args.train_data_ratio)
train, _ = split_dataset_random(dataset, train_data_size, args.seed)
# Set up the predictor.
predictor = set_up_predictor(args.method, args.unit_num, args.conv_layers,
class_num)
# Set up the iterator.
train_iter = SerialIterator(train, args.batchsize)
# Set up the regressor.
metrics_fun = {
'mean_abs_error': MeanAbsError(scaler=scaler),
'root_mean_sqr_error': RootMeanSqrError(scaler=scaler)
}
regressor = Regressor(predictor,
lossfun=F.mean_squared_error,
metrics_fun=metrics_fun,
device=args.gpu)
# Set up the optimizer.
optimizer = optimizers.Adam()
optimizer.setup(regressor)
# Set up the updater.
updater = training.StandardUpdater(train_iter,
optimizer,
device=args.gpu,
converter=concat_mols)
# Set up the trainer.
print('Training...')
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)
trainer.extend(E.snapshot(), trigger=(args.epoch, 'epoch'))
trainer.extend(E.LogReport())
trainer.extend(
E.PrintReport([
'epoch', 'main/loss', 'main/mean_abs_error',
'main/root_mean_sqr_error', 'elapsed_time'
]))
trainer.extend(E.ProgressBar())
trainer.run()
# Save the regressor's parameters.
model_path = os.path.join(args.out, args.model_filename)
print('Saving the trained model to {}...'.format(model_path))
regressor.save_pickle(model_path, protocol=args.protocol)
# Save the standard scaler's parameters.
if scaler is not None:
with open(os.path.join(args.out, 'scaler.pkl'), mode='wb') as f:
pickle.dump(scaler, f, protocol=args.protocol)
def test_get_datasets(self, data):
dataset = NumpyTupleDataset(*data)
datasets = dataset.get_datasets()
assert len(datasets) == len(data)
for i in range(len(datasets)):
numpy.testing.assert_array_equal(datasets[i], data[i])
def main():
# Parse the arguments.
args = parse_arguments()
# Set up some useful variables that will be used later on.
method = args.method
if args.label != 'all':
labels = args.label
cache_dir = os.path.join('input', '{}_{}'.format(method, labels))
class_num = len(labels) if isinstance(labels, list) else 1
else:
labels = None
cache_dir = os.path.join('input', '{}_all'.format(method))
class_num = len(D.get_qm9_label_names())
# Get the filename corresponding to the cached dataset, based on the amount
# of data samples that need to be parsed from the original dataset.
num_data = args.num_data
if num_data >= 0:
dataset_filename = 'data_{}.npz'.format(num_data)
else:
dataset_filename = 'data.npz'
# Load the cached dataset.
dataset_cache_path = os.path.join(cache_dir, dataset_filename)
dataset = None
if os.path.exists(dataset_cache_path):
print('Loading cached dataset from {}.'.format(dataset_cache_path))
dataset = NumpyTupleDataset.load(dataset_cache_path)
if dataset is None:
print('Preprocessing dataset...')
preprocessor = preprocess_method_dict[method]()
if num_data >= 0:
# Select the first `num_data` samples from the dataset.
target_index = numpy.arange(num_data)
dataset = D.get_qm9(preprocessor, labels=labels,
target_index=target_index)
else:
# Load the entire dataset.
dataset = D.get_qm9(preprocessor, labels=labels)
# Cache the laded dataset.
if not os.path.exists(cache_dir):
os.makedirs(cache_dir)
NumpyTupleDataset.save(dataset_cache_path, dataset)
# Scale the label values, if necessary.
if args.scale == 'standardize':
print('Applying standard scaling to the labels.')
scaler = StandardScaler()
scaled_t = scaler.fit_transform(dataset.get_datasets()[-1])
dataset = NumpyTupleDataset(*(dataset.get_datasets()[:-1]
+ (scaled_t,)))
else:
print('No standard scaling was selected.')
scaler = None
# Split the dataset into training and validation.
train_data_size = int(len(dataset) * args.train_data_ratio)
train, valid = split_dataset_random(dataset, train_data_size, args.seed)
# Set up the predictor.
predictor = set_up_predictor(method, args.unit_num, args.conv_layers,
class_num, scaler)
# Set up the iterators.
train_iter = iterators.SerialIterator(train, args.batchsize)
valid_iter = iterators.SerialIterator(valid, args.batchsize, repeat=False,
shuffle=False)
# Set up the regressor.
device = args.gpu
metrics_fun = {'mae': MeanAbsError(scaler=scaler),
'rmse': RootMeanSqrError(scaler=scaler)}
regressor = Regressor(predictor, lossfun=F.mean_squared_error,
metrics_fun=metrics_fun, device=device)
# Set up the optimizer.
optimizer = optimizers.Adam()
optimizer.setup(regressor)
# Set up the updater.
updater = training.StandardUpdater(train_iter, optimizer, device=device,
converter=concat_mols)
# Set up the trainer.
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)
trainer.extend(E.Evaluator(valid_iter, regressor, device=device,
converter=concat_mols))
trainer.extend(E.snapshot(), trigger=(args.epoch, 'epoch'))
trainer.extend(E.LogReport())
trainer.extend(E.PrintReport([
'epoch', 'main/loss', 'main/mae', 'main/rmse', 'validation/main/loss',
'validation/main/mae', 'validation/main/rmse', 'elapsed_time']))
trainer.extend(E.ProgressBar())
trainer.run()
# Save the regressor's parameters.
model_path = os.path.join(args.out, args.model_filename)
print('Saving the trained model to {}...'.format(model_path))
regressor.save_pickle(model_path, protocol=args.protocol)
def main():
# Supported preprocessing/network list
method_list = ['nfp', 'ggnn', 'schnet', 'weavenet', 'rsgcn']
label_names = [
'A', 'B', 'C', 'mu', 'alpha', 'h**o', 'lumo', 'gap', 'r2', 'zpve',
'U0', 'U', 'H', 'G', 'Cv'
]
scale_list = ['standardize', 'none']
parser = argparse.ArgumentParser(description='Regression with QM9.')
parser.add_argument('--method',
'-m',
type=str,
choices=method_list,
default='nfp')
parser.add_argument('--label',
'-l',
type=str,
choices=label_names,
default='',
help='target label for regression, '
'empty string means to predict all '
'property at once')
parser.add_argument('--scale',
type=str,
choices=scale_list,
default='standardize',
help='Label scaling method')
parser.add_argument('--conv-layers', '-c', type=int, default=4)
parser.add_argument('--batchsize', '-b', type=int, default=32)
parser.add_argument('--gpu', '-g', type=int, default=-1)
parser.add_argument('--out', '-o', type=str, default='result')
parser.add_argument('--epoch', '-e', type=int, default=20)
parser.add_argument('--unit-num', '-u', type=int, default=16)
parser.add_argument('--seed', '-s', type=int, default=777)
parser.add_argument('--train-data-ratio', '-t', type=float, default=0.7)
args = parser.parse_args()
seed = args.seed
train_data_ratio = args.train_data_ratio
method = args.method
if args.label:
labels = args.label
cache_dir = os.path.join('input', '{}_{}'.format(method, labels))
class_num = len(labels) if isinstance(labels, list) else 1
else:
labels = None
cache_dir = os.path.join('input', '{}_all'.format(method))
class_num = len(D.get_qm9_label_names())
# Dataset preparation
dataset = None
if os.path.exists(cache_dir):
print('load from cache {}'.format(cache_dir))
dataset = NumpyTupleDataset.load(os.path.join(cache_dir, 'data.npz'))
if dataset is None:
print('preprocessing dataset...')
preprocessor = preprocess_method_dict[method]()
dataset = D.get_qm9(preprocessor, labels=labels)
os.makedirs(cache_dir)
NumpyTupleDataset.save(os.path.join(cache_dir, 'data.npz'), dataset)
if args.scale == 'standardize':
# Standard Scaler for labels
ss = StandardScaler()
labels = ss.fit_transform(dataset.get_datasets()[-1])
dataset = NumpyTupleDataset(*dataset.get_datasets()[:-1], labels)
train_data_size = int(len(dataset) * train_data_ratio)
train, val = split_dataset_random(dataset, train_data_size, seed)
# Network
n_unit = args.unit_num
conv_layers = args.conv_layers
if method == 'nfp':
print('Train NFP model...')
model = GraphConvPredictor(
NFP(out_dim=n_unit, hidden_dim=n_unit, n_layers=conv_layers),
MLP(out_dim=class_num, hidden_dim=n_unit))
elif method == 'ggnn':
print('Train GGNN model...')
model = GraphConvPredictor(
GGNN(out_dim=n_unit, hidden_dim=n_unit, n_layers=conv_layers),
MLP(out_dim=class_num, hidden_dim=n_unit))
elif method == 'schnet':
print('Train SchNet model...')
model = GraphConvPredictor(
SchNet(out_dim=class_num, hidden_dim=n_unit, n_layers=conv_layers),
None)
elif method == 'weavenet':
print('Train WeaveNet model...')
n_atom = 20
n_sub_layer = 1
weave_channels = [50] * conv_layers
model = GraphConvPredictor(
WeaveNet(weave_channels=weave_channels,
hidden_dim=n_unit,
n_sub_layer=n_sub_layer,
n_atom=n_atom), MLP(out_dim=class_num, hidden_dim=n_unit))
elif method == 'rsgcn':
print('Train RSGCN model...')
model = GraphConvPredictor(
RSGCN(out_dim=n_unit, hidden_dim=n_unit, n_layers=conv_layers),
MLP(out_dim=class_num, hidden_dim=n_unit))
else:
raise ValueError('[ERROR] Invalid method {}'.format(method))
train_iter = I.SerialIterator(train, args.batchsize)
val_iter = I.SerialIterator(val,
args.batchsize,
repeat=False,
shuffle=False)
def scaled_abs_error(x0, x1):
if isinstance(x0, Variable):
x0 = cuda.to_cpu(x0.data)
if isinstance(x1, Variable):
x1 = cuda.to_cpu(x1.data)
if args.scale == 'standardize':
scaled_x0 = ss.inverse_transform(cuda.to_cpu(x0))
scaled_x1 = ss.inverse_transform(cuda.to_cpu(x1))
diff = scaled_x0 - scaled_x1
elif args.scale == 'none':
diff = cuda.to_cpu(x0) - cuda.to_cpu(x1)
return numpy.mean(numpy.absolute(diff), axis=0)[0]
classifier = L.Classifier(model,
lossfun=F.mean_squared_error,
accfun=scaled_abs_error)
if args.gpu >= 0:
chainer.cuda.get_device_from_id(args.gpu).use()
classifier.to_gpu()
optimizer = O.Adam()
optimizer.setup(classifier)
updater = training.StandardUpdater(train_iter,
optimizer,
device=args.gpu,
converter=concat_mols)
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)
trainer.extend(
E.Evaluator(val_iter,
classifier,
device=args.gpu,
converter=concat_mols))
trainer.extend(E.snapshot(), trigger=(args.epoch, 'epoch'))
trainer.extend(E.LogReport())
trainer.extend(
E.PrintReport([
'epoch', 'main/loss', 'main/accuracy', 'validation/main/loss',
'validation/main/accuracy', 'elapsed_time'
]))
trainer.extend(E.ProgressBar())
trainer.run()
def main():
# Supported preprocessing/network list
method_list = ['nfp', 'ggnn', 'schnet', 'weavenet', 'rsgcn']
scale_list = ['standardize', 'none']
parser = argparse.ArgumentParser(
description='Regression with own dataset.')
parser.add_argument('--datafile', type=str, default='dataset.csv')
parser.add_argument('--method',
'-m',
type=str,
choices=method_list,
default='nfp')
parser.add_argument('--label',
'-l',
nargs='+',
default=['value1', 'value2'],
help='target label for regression')
parser.add_argument('--scale',
type=str,
choices=scale_list,
default='standardize',
help='Label scaling method')
parser.add_argument('--conv-layers', '-c', type=int, default=4)
parser.add_argument('--batchsize', '-b', type=int, default=32)
parser.add_argument('--gpu', '-g', type=int, default=-1)
parser.add_argument('--out', '-o', type=str, default='result')
parser.add_argument('--epoch', '-e', type=int, default=20)
parser.add_argument('--unit-num', '-u', type=int, default=16)
parser.add_argument('--seed', '-s', type=int, default=777)
parser.add_argument('--train-data-ratio', '-t', type=float, default=0.7)
parser.add_argument('--protocol', type=int, default=2)
args = parser.parse_args()
seed = args.seed
train_data_ratio = args.train_data_ratio
method = args.method
if args.label:
labels = args.label
class_num = len(labels) if isinstance(labels, list) else 1
else:
sys.exit("Error: No target label is specified.")
# Dataset preparation
# Postprocess is required for regression task
def postprocess_label(label_list):
return numpy.asarray(label_list, dtype=numpy.float32)
print('Preprocessing dataset...')
preprocessor = preprocess_method_dict[method]()
parser = CSVFileParser(preprocessor,
postprocess_label=postprocess_label,
labels=labels,
smiles_col='SMILES')
dataset = parser.parse(args.datafile)["dataset"]
if args.scale == 'standardize':
# Standard Scaler for labels
scaler = StandardScaler()
labels = scaler.fit_transform(dataset.get_datasets()[-1])
dataset = NumpyTupleDataset(*(dataset.get_datasets()[:-1] +
(labels, )))
else:
# Not use scaler
scaler = None
train_data_size = int(len(dataset) * train_data_ratio)
train, val = split_dataset_random(dataset, train_data_size, seed)
# Network
n_unit = args.unit_num
conv_layers = args.conv_layers
if method == 'nfp':
print('Train NFP model...')
model = GraphConvPredictor(
NFP(out_dim=n_unit, hidden_dim=n_unit, n_layers=conv_layers),
MLP(out_dim=class_num, hidden_dim=n_unit))
elif method == 'ggnn':
print('Train GGNN model...')
model = GraphConvPredictor(
GGNN(out_dim=n_unit, hidden_dim=n_unit, n_layers=conv_layers),
MLP(out_dim=class_num, hidden_dim=n_unit))
elif method == 'schnet':
print('Train SchNet model...')
model = GraphConvPredictor(
SchNet(out_dim=class_num, hidden_dim=n_unit, n_layers=conv_layers),
None)
elif method == 'weavenet':
print('Train WeaveNet model...')
n_atom = 20
n_sub_layer = 1
weave_channels = [50] * conv_layers
model = GraphConvPredictor(
WeaveNet(weave_channels=weave_channels,
hidden_dim=n_unit,
n_sub_layer=n_sub_layer,
n_atom=n_atom), MLP(out_dim=class_num, hidden_dim=n_unit))
elif method == 'rsgcn':
print('Train RSGCN model...')
model = GraphConvPredictor(
RSGCN(out_dim=n_unit, hidden_dim=n_unit, n_layers=conv_layers),
MLP(out_dim=class_num, hidden_dim=n_unit))
else:
raise ValueError('[ERROR] Invalid method {}'.format(method))
train_iter = iterators.SerialIterator(train, args.batchsize)
val_iter = iterators.SerialIterator(val,
args.batchsize,
repeat=False,
shuffle=False)
regressor = Regressor(
model,
lossfun=F.mean_squared_error,
metrics_fun={'abs_error': ScaledAbsError(scaler=scaler)},
device=args.gpu)
optimizer = optimizers.Adam()
optimizer.setup(regressor)
updater = training.StandardUpdater(train_iter,
optimizer,
device=args.gpu,
converter=concat_mols)
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)
trainer.extend(
E.Evaluator(val_iter,
regressor,
device=args.gpu,
converter=concat_mols))
trainer.extend(E.snapshot(), trigger=(args.epoch, 'epoch'))
trainer.extend(E.LogReport())
# Note that original scale absolute errors are reported in
# (validation/)main/abs_error
trainer.extend(
E.PrintReport([
'epoch', 'main/loss', 'main/abs_error', 'validation/main/loss',
'validation/main/abs_error', 'elapsed_time'
]))
trainer.extend(E.ProgressBar())
trainer.run()
# --- save regressor's parameters ---
protocol = args.protocol
model_path = os.path.join(args.out, 'model.npz')
print('saving trained model to {}'.format(model_path))
serializers.save_npz(model_path, regressor)
if scaler is not None:
with open(os.path.join(args.out, 'scaler.pkl'), mode='wb') as f:
pickle.dump(scaler, f, protocol=protocol)
# Example of prediction using trained model
smiles = 'c1ccccc1'
mol = Chem.MolFromSmiles(smiles)
preprocessor = preprocess_method_dict[method]()
standardized_smiles, mol = preprocessor.prepare_smiles_and_mol(mol)
input_features = preprocessor.get_input_features(mol)
atoms, adjs = concat_mols([input_features], device=args.gpu)
prediction = model(atoms, adjs).data[0]
if scaler is not None:
prediction = scaler.inverse_transform(prediction)
print('Prediction for {}:'.format(smiles))
for i, label in enumerate(args.label):
print('{}: {}'.format(label, prediction[i]))
def main():
# Parse the arguments.
args = parse_arguments()
device = args.gpu
# Set up some useful variables that will be used later on.
method = args.method
if args.label != 'all':
label = args.label
cache_dir = os.path.join('input', '{}_{}'.format(method, label))
labels = [label]
else:
labels = D.get_qm9_label_names()
cache_dir = os.path.join('input', '{}_all'.format(method))
# Get the filename corresponding to the cached dataset, based on the amount
# of data samples that need to be parsed from the original dataset.
num_data = args.num_data
if num_data >= 0:
dataset_filename = 'data_{}.npz'.format(num_data)
else:
dataset_filename = 'data.npz'
# Load the cached dataset.
dataset_cache_path = os.path.join(cache_dir, dataset_filename)
dataset = None
if os.path.exists(dataset_cache_path):
print('Loading cached data from {}.'.format(dataset_cache_path))
dataset = NumpyTupleDataset.load(dataset_cache_path)
if dataset is None:
print('Preprocessing dataset...')
preprocessor = preprocess_method_dict[method]()
dataset = D.get_qm9(preprocessor, labels=labels)
# Cache the newly preprocessed dataset.
if not os.path.exists(cache_dir):
os.mkdir(cache_dir)
NumpyTupleDataset.save(dataset_cache_path, dataset)
# Use a predictor with scaled output labels.
model_path = os.path.join(args.in_dir, args.model_filename)
regressor = Regressor.load_pickle(model_path, device=device)
scaler = regressor.predictor.scaler
if scaler is not None:
original_t = dataset.get_datasets()[-1]
if args.gpu >= 0:
scaled_t = cuda.to_cpu(scaler.transform(
cuda.to_gpu(original_t)))
else:
scaled_t = scaler.transform(original_t)
dataset = NumpyTupleDataset(*(dataset.get_datasets()[:-1] +
(scaled_t,)))
# Split the dataset into training and testing.
train_data_size = int(len(dataset) * args.train_data_ratio)
_, test = split_dataset_random(dataset, train_data_size, args.seed)
# This callback function extracts only the inputs and discards the labels.
def extract_inputs(batch, device=None):
return concat_mols(batch, device=device)[:-1]
def postprocess_fn(x):
if scaler is not None:
scaled_x = scaler.inverse_transform(x)
return scaled_x
else:
return x
# Predict the output labels.
print('Predicting...')
y_pred = regressor.predict(
test, converter=extract_inputs,
postprocess_fn=postprocess_fn)
# Extract the ground-truth labels.
t = concat_mols(test, device=device)[-1]
original_t = cuda.to_cpu(scaler.inverse_transform(t))
# Construct dataframe.
df_dict = {}
for i, l in enumerate(labels):
df_dict.update({'y_pred_{}'.format(l): y_pred[:, i],
't_{}'.format(l): original_t[:, i], })
df = pandas.DataFrame(df_dict)
# Show a prediction/ground truth table with 5 random examples.
print(df.sample(5))
n_eval = 10
for target_label in range(y_pred.shape[1]):
label_name = labels[target_label]
diff = y_pred[:n_eval, target_label] - original_t[:n_eval,
target_label]
print('label_name = {}, y_pred = {}, t = {}, diff = {}'
.format(label_name, y_pred[:n_eval, target_label],
original_t[:n_eval, target_label], diff))
# Run an evaluator on the test dataset.
print('Evaluating...')
test_iterator = SerialIterator(test, 16, repeat=False, shuffle=False)
eval_result = Evaluator(test_iterator, regressor, converter=concat_mols,
device=device)()
print('Evaluation result: ', eval_result)
# Save the evaluation results.
save_json(os.path.join(args.in_dir, 'eval_result.json'), eval_result)
# Calculate mean abs error for each label
mae = numpy.mean(numpy.abs(y_pred - original_t), axis=0)
eval_result = {}
for i, l in enumerate(labels):
eval_result.update({l: mae[i]})
save_json(os.path.join(args.in_dir, 'eval_result_mae.json'), eval_result)
def main():
# Supported preprocessing/network list
method_list = ['nfp', 'ggnn', 'schnet', 'weavenet', 'rsgcn']
label_names = [
'A', 'B', 'C', 'mu', 'alpha', 'h**o', 'lumo', 'gap', 'r2', 'zpve',
'U0', 'U', 'H', 'G', 'Cv'
]
scale_list = ['standardize', 'none']
parser = argparse.ArgumentParser(description='Regression with QM9.')
parser.add_argument('--method',
'-m',
type=str,
choices=method_list,
default='nfp')
parser.add_argument('--label',
'-l',
type=str,
choices=label_names,
default='',
help='target label for regression, '
'empty string means to predict all '
'property at once')
parser.add_argument('--scale',
type=str,
choices=scale_list,
default='standardize',
help='Label scaling method')
parser.add_argument('--conv-layers', '-c', type=int, default=4)
parser.add_argument('--batchsize', '-b', type=int, default=32)
parser.add_argument('--gpu', '-g', type=int, default=-1)
parser.add_argument('--out', '-o', type=str, default='result')
parser.add_argument('--epoch', '-e', type=int, default=20)
parser.add_argument('--unit-num', '-u', type=int, default=16)
parser.add_argument('--seed', '-s', type=int, default=777)
parser.add_argument('--train-data-ratio', '-t', type=float, default=0.7)
parser.add_argument('--protocol', type=int, default=2)
parser.add_argument('--model-filename', type=str, default='regressor.pkl')
parser.add_argument('--num-data',
type=int,
default=-1,
help='Number of data to be parsed from parser.'
'-1 indicates to parse all data.')
args = parser.parse_args()
seed = args.seed
train_data_ratio = args.train_data_ratio
method = args.method
if args.label:
labels = args.label
cache_dir = os.path.join('input', '{}_{}'.format(method, labels))
class_num = len(labels) if isinstance(labels, list) else 1
else:
labels = None
cache_dir = os.path.join('input', '{}_all'.format(method))
class_num = len(D.get_qm9_label_names())
# Dataset preparation
dataset = None
num_data = args.num_data
if num_data >= 0:
dataset_filename = 'data_{}.npz'.format(num_data)
else:
dataset_filename = 'data.npz'
dataset_cache_path = os.path.join(cache_dir, dataset_filename)
if os.path.exists(dataset_cache_path):
print('load from cache {}'.format(dataset_cache_path))
dataset = NumpyTupleDataset.load(dataset_cache_path)
if dataset is None:
print('preprocessing dataset...')
preprocessor = preprocess_method_dict[method]()
if num_data >= 0:
# only use first 100 for debug
target_index = numpy.arange(num_data)
dataset = D.get_qm9(preprocessor,
labels=labels,
target_index=target_index)
else:
dataset = D.get_qm9(preprocessor, labels=labels)
os.makedirs(cache_dir)
NumpyTupleDataset.save(dataset_cache_path, dataset)
if args.scale == 'standardize':
# Standard Scaler for labels
ss = StandardScaler()
labels = ss.fit_transform(dataset.get_datasets()[-1])
else:
ss = None
dataset = NumpyTupleDataset(*(dataset.get_datasets()[:-1] + (labels, )))
train_data_size = int(len(dataset) * train_data_ratio)
train, val = split_dataset_random(dataset, train_data_size, seed)
# Network
n_unit = args.unit_num
conv_layers = args.conv_layers
if method == 'nfp':
print('Train NFP model...')
model = GraphConvPredictor(
NFP(out_dim=n_unit, hidden_dim=n_unit, n_layers=conv_layers),
MLP(out_dim=class_num, hidden_dim=n_unit))
elif method == 'ggnn':
print('Train GGNN model...')
model = GraphConvPredictor(
GGNN(out_dim=n_unit, hidden_dim=n_unit, n_layers=conv_layers),
MLP(out_dim=class_num, hidden_dim=n_unit))
elif method == 'schnet':
print('Train SchNet model...')
model = GraphConvPredictor(
SchNet(out_dim=class_num, hidden_dim=n_unit, n_layers=conv_layers),
None)
elif method == 'weavenet':
print('Train WeaveNet model...')
n_atom = 20
n_sub_layer = 1
weave_channels = [50] * conv_layers
model = GraphConvPredictor(
WeaveNet(weave_channels=weave_channels,
hidden_dim=n_unit,
n_sub_layer=n_sub_layer,
n_atom=n_atom), MLP(out_dim=class_num, hidden_dim=n_unit))
elif method == 'rsgcn':
print('Train RSGCN model...')
model = GraphConvPredictor(
RSGCN(out_dim=n_unit, hidden_dim=n_unit, n_layers=conv_layers),
MLP(out_dim=class_num, hidden_dim=n_unit))
else:
raise ValueError('[ERROR] Invalid method {}'.format(method))
train_iter = I.SerialIterator(train, args.batchsize)
val_iter = I.SerialIterator(val,
args.batchsize,
repeat=False,
shuffle=False)
regressor = Regressor(
model,
lossfun=F.mean_squared_error,
metrics_fun={'abs_error': ScaledAbsError(scale=args.scale, ss=ss)},
device=args.gpu)
optimizer = O.Adam()
optimizer.setup(regressor)
updater = training.StandardUpdater(train_iter,
optimizer,
device=args.gpu,
converter=concat_mols)
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)
trainer.extend(
E.Evaluator(val_iter,
regressor,
device=args.gpu,
converter=concat_mols))
trainer.extend(E.snapshot(), trigger=(args.epoch, 'epoch'))
trainer.extend(E.LogReport())
trainer.extend(
E.PrintReport([
'epoch', 'main/loss', 'main/abs_error', 'validation/main/loss',
'validation/main/abs_error', 'elapsed_time'
]))
trainer.extend(E.ProgressBar())
trainer.run()
# --- save regressor & standardscaler ---
protocol = args.protocol
regressor.save_pickle(os.path.join(args.out, args.model_filename),
protocol=protocol)
if args.scale == 'standardize':
with open(os.path.join(args.out, 'ss.pkl'), mode='wb') as f:
pickle.dump(ss, f, protocol=protocol)
def main():
# Parse the arguments.
args = parse_arguments()
augment = False if args.augment == 'False' else True
multi_gpu = False if args.multi_gpu == 'False' else True
if args.label:
labels = args.label
class_num = len(labels) if isinstance(labels, list) else 1
else:
raise ValueError('No target label was specified.')
# Dataset preparation. Postprocessing is required for the regression task.
def postprocess_label(label_list):
label_arr = np.asarray(label_list, dtype=np.int32)
return label_arr
# Apply a preprocessor to the dataset.
logging.info('Preprocess train dataset and test dataset...')
preprocessor = preprocess_method_dict[args.method]()
parser = CSVFileParserForPair(preprocessor,
postprocess_label=postprocess_label,
labels=labels,
smiles_cols=['smiles_1', 'smiles_2'])
train = parser.parse(args.train_datafile)['dataset']
test = parser.parse(args.valid_datafile)['dataset']
if augment:
logging.info('Utilizing data augmentation in train set')
train = augment_dataset(train)
num_train = train.get_datasets()[0].shape[0]
num_test = test.get_datasets()[0].shape[0]
logging.info('Train/test split: {}/{}'.format(num_train, num_test))
if len(args.net_hidden_dims):
net_hidden_dims = tuple([
int(net_hidden_dim)
for net_hidden_dim in args.net_hidden_dims.split(',')
])
else:
net_hidden_dims = ()
weight_tying = False if args.weight_tying == 'False' else True
fp_batch_normalization = True if args.fp_bn == 'True' else False
predictor = set_up_predictor(
method=args.method,
fp_hidden_dim=args.fp_hidden_dim,
fp_out_dim=args.fp_out_dim,
conv_layers=args.conv_layers,
concat_hidden=args.concat_hidden,
fp_dropout_rate=args.fp_dropout_rate,
fp_batch_normalization=fp_batch_normalization,
net_hidden_dims=net_hidden_dims,
class_num=class_num,
sim_method=args.sim_method,
weight_typing=weight_tying,
symmetric=args.symmetric,
attn_model=args.attn,
)
if args.train_pos_neg_ratio != -1.:
# Set up the iterator.
train_dataset = train.get_datasets()
atoms1_train, adjs1_train, atoms2_train, adjs2_train, labels_train = train_dataset
labels_train = np.squeeze(labels_train)
train_dataset_arr = np.concatenate([
item[:, None] if len(item.shape) == 1 else item
for item in list(train_dataset)
],
axis=1)
pos_train_dataset_arr = train_dataset_arr[labels_train == 1]
num_pos_train = pos_train_dataset_arr.shape[0]
pos_train_indices = np.arange(0, num_pos_train)
neg_train_dataset_arr = train_dataset_arr[labels_train == 0]
num_neg_train = neg_train_dataset_arr.shape[0]
pos_neg_train_ratio = args.train_pos_neg_ratio
num_pos_train = int(pos_neg_train_ratio * num_neg_train)
np.random.seed(777)
np.random.shuffle(pos_train_indices)
pos_train_indices = pos_train_indices[:num_pos_train]
pos_train_dataset_arr = pos_train_dataset_arr[pos_train_indices]
new_train_dataset_arr = np.concatenate(
(pos_train_dataset_arr, neg_train_dataset_arr), axis=0)
atoms1_train, adjs1_train = new_train_dataset_arr[:,
0], new_train_dataset_arr[:,
1]
atoms2_train, adjs2_train = new_train_dataset_arr[:,
2], new_train_dataset_arr[:,
3]
labels_train = new_train_dataset_arr[:, 4].astype(np.int32)
labels_train = np.expand_dims(labels_train, axis=1)
train = NumpyTupleDataset(atoms1_train, adjs1_train, atoms2_train,
adjs2_train, labels_train)
num_train = train.get_datasets()[0].shape[0]
num_test = test.get_datasets()[0].shape[0]
logging.info('Train pos-neg ratio is {:.4f}'.format(
args.train_pos_neg_ratio))
logging.info('Train/test number is {}/{}'.format(num_train, num_test))
# if args.loss_func == 'hinge':
# modify_dataset_for_hinge(train)
# Set up the iterator.
train_iter = SerialIterator(train, args.batchsize)
test_iter = SerialIterator(test,
args.batchsize,
repeat=False,
shuffle=False)
metrics_fun = {'accuracy': F.binary_accuracy}
loss_func = F.sigmoid_cross_entropy
if args.loss_func == 'hinge':
logging.info('Loss function is {}'.format(args.loss_func))
loss_func = F.hinge
metrics_fun = {'accuracy': F.accuracy}
classifier = Classifier(predictor,
lossfun=loss_func,
metrics_fun=metrics_fun,
device=args.gpu)
# Set up the optimizer.
optimizer = optimizers.Adam(alpha=args.learning_rate,
weight_decay_rate=args.weight_decay_rate)
# optimizer = optimizers.Adam()
# optimizer = optimizers.SGD(lr=args.learning_rate)
optimizer.setup(classifier)
# add regularization
if args.max_norm > 0:
optimizer.add_hook(
chainer.optimizer.GradientClipping(threshold=args.max_norm))
if args.l2_rate > 0:
optimizer.add_hook(chainer.optimizer.WeightDecay(rate=args.l2_rate))
if args.l1_rate > 0:
optimizer.add_hook(chainer.optimizer.Lasso(rate=args.l1_rate))
# Set up the updater.
if multi_gpu:
logging.info('Using multiple GPUs')
updater = training.ParallelUpdater(train_iter,
optimizer,
devices={
'main': 0,
'second': 1
},
converter=concat_mols)
else:
logging.info('Using single GPU')
updater = training.StandardUpdater(train_iter,
optimizer,
device=args.gpu,
converter=concat_mols)
# Set up the trainer.
logging.info('Training...')
# add stop_trigger parameter
early_stop = triggers.EarlyStoppingTrigger(monitor='validation/main/loss',
patients=50,
max_trigger=(500, 'epoch'))
out = 'output' + '/' + args.out
trainer = training.Trainer(updater, stop_trigger=early_stop, out=out)
# trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)
trainer.extend(
E.Evaluator(test_iter,
classifier,
device=args.gpu,
converter=concat_mols))
train_eval_iter = SerialIterator(train,
args.batchsize,
repeat=False,
shuffle=False)
trainer.extend(
AccuracyEvaluator(train_eval_iter,
classifier,
eval_func=predictor,
device=args.gpu,
converter=concat_mols,
name='train_acc',
pos_labels=1,
ignore_labels=-1,
raise_value_error=False))
# extension name='validation' is already used by `Evaluator`,
# instead extension name `val` is used.
trainer.extend(
AccuracyEvaluator(test_iter,
classifier,
eval_func=predictor,
device=args.gpu,
converter=concat_mols,
name='val_acc',
pos_labels=1,
ignore_labels=-1))
trainer.extend(
ROCAUCEvaluator(train_eval_iter,
classifier,
eval_func=predictor,
device=args.gpu,
converter=concat_mols,
name='train_roc',
pos_labels=1,
ignore_labels=-1,
raise_value_error=False))
# extension name='validation' is already used by `Evaluator`,
# instead extension name `val` is used.
trainer.extend(
ROCAUCEvaluator(test_iter,
classifier,
eval_func=predictor,
device=args.gpu,
converter=concat_mols,
name='val_roc',
pos_labels=1,
ignore_labels=-1))
trainer.extend(
PRCAUCEvaluator(train_eval_iter,
classifier,
eval_func=predictor,
device=args.gpu,
converter=concat_mols,
name='train_prc',
pos_labels=1,
ignore_labels=-1,
raise_value_error=False))
# extension name='validation' is already used by `Evaluator`,
# instead extension name `val` is used.
trainer.extend(
PRCAUCEvaluator(test_iter,
classifier,
eval_func=predictor,
device=args.gpu,
converter=concat_mols,
name='val_prc',
pos_labels=1,
ignore_labels=-1))
# trainer.extend(PrecisionEvaluator(
# train_eval_iter, classifier, eval_func=predictor,
# device=args.gpu, converter=concat_mols, name='train_p',
# pos_labels=1, ignore_labels=-1, raise_value_error=False))
# # extension name='validation' is already used by `Evaluator`,
# # instead extension name `val` is used.
# trainer.extend(PrecisionEvaluator(
# val_iter, classifier, eval_func=predictor,
# device=args.gpu, converter=concat_mols, name='val_p',
# pos_labels=1, ignore_labels=-1))
#
# trainer.extend(RecallEvaluator(
# train_eval_iter, classifier, eval_func=predictor,
# device=args.gpu, converter=concat_mols, name='train_r',
# pos_labels=1, ignore_labels=-1, raise_value_error=False))
# # extension name='validation' is already used by `Evaluator`,
# # instead extension name `val` is used.
# trainer.extend(RecallEvaluator(
# val_iter, classifier, eval_func=predictor,
# device=args.gpu, converter=concat_mols, name='val_r',
# pos_labels=1, ignore_labels=-1))
trainer.extend(
F1Evaluator(train_eval_iter,
classifier,
eval_func=predictor,
device=args.gpu,
converter=concat_mols,
name='train_f',
pos_labels=1,
ignore_labels=-1,
raise_value_error=False))
# extension name='validation' is already used by `Evaluator`,
# instead extension name `val` is used.
trainer.extend(
F1Evaluator(test_iter,
classifier,
eval_func=predictor,
device=args.gpu,
converter=concat_mols,
name='val_f',
pos_labels=1,
ignore_labels=-1))
# apply shift strategy to learning rate every 10 epochs
# trainer.extend(E.ExponentialShift('alpha', args.exp_shift_rate), trigger=(10, 'epoch'))
if args.exp_shift_strategy == 1:
trainer.extend(E.ExponentialShift('alpha', args.exp_shift_rate),
trigger=triggers.ManualScheduleTrigger(
[10, 20, 30, 40, 50, 60], 'epoch'))
elif args.exp_shift_strategy == 2:
trainer.extend(E.ExponentialShift('alpha', args.exp_shift_rate),
trigger=triggers.ManualScheduleTrigger(
[5, 10, 15, 20, 25, 30], 'epoch'))
elif args.exp_shift_strategy == 3:
trainer.extend(E.ExponentialShift('alpha', args.exp_shift_rate),
trigger=triggers.ManualScheduleTrigger(
[5, 10, 15, 20, 25, 30, 40, 50, 60, 70], 'epoch'))
else:
raise ValueError('No such strategy to adapt learning rate')
# # observation of learning rate
trainer.extend(E.observe_lr(), trigger=(1, 'iteration'))
entries = [
'epoch',
'main/loss',
'train_acc/main/accuracy',
'train_roc/main/roc_auc',
'train_prc/main/prc_auc',
# 'train_p/main/precision', 'train_r/main/recall',
'train_f/main/f1',
'validation/main/loss',
'val_acc/main/accuracy',
'val_roc/main/roc_auc',
'val_prc/main/prc_auc',
# 'val_p/main/precision', 'val_r/main/recall',
'val_f/main/f1',
'lr',
'elapsed_time'
]
trainer.extend(E.PrintReport(entries=entries))
# change from 10 to 2 on Mar. 1 2019
trainer.extend(E.snapshot(), trigger=(2, 'epoch'))
trainer.extend(E.LogReport())
trainer.extend(E.ProgressBar())
trainer.extend(
E.PlotReport(['main/loss', 'validation/main/loss'],
'epoch',
file_name='loss.png'))
trainer.extend(
E.PlotReport(['train_acc/main/accuracy', 'val_acc/main/accuracy'],
'epoch',
file_name='accuracy.png'))
if args.resume:
resume_path = os.path.join(out, args.resume)
logging.info(
'Resume training according to snapshot in {}'.format(resume_path))
chainer.serializers.load_npz(resume_path, trainer)
trainer.run()
# Save the regressor's parameters.
model_path = os.path.join(out, args.model_filename)
logging.info('Saving the trained models to {}...'.format(model_path))
classifier.save_pickle(model_path, protocol=args.protocol)
def train(gpu, method, epoch, batchsize, n_unit, conv_layers, dataset, smiles,
M, n_split, split_idx, order):
n = len(dataset)
assert len(order) == n
left_idx = (n // n_split) * split_idx
is_right_most_split = (n_split == split_idx + 1)
if is_right_most_split:
test_order = order[left_idx:]
train_order = order[:left_idx]
else:
right_idx = (n // n_split) * (split_idx + 1)
test_order = order[left_idx:right_idx]
train_order = np.concatenate([order[:left_idx], order[right_idx:]])
new_order = np.concatenate([train_order, test_order])
n_train = len(train_order)
# Standard Scaler for labels
ss = StandardScaler()
labels = dataset.get_datasets()[-1]
train_label = labels[new_order[:n_train]]
ss = ss.fit(train_label) # fit only by train
labels = ss.transform(dataset.get_datasets()[-1])
dataset = NumpyTupleDataset(*(dataset.get_datasets()[:-1] + (labels, )))
dataset_train = SubDataset(dataset, 0, n_train, new_order)
dataset_test = SubDataset(dataset, n_train, n, new_order)
# Network
model = predictor.build_predictor(method,
n_unit,
conv_layers,
1,
dropout_ratio=0.25,
n_layers=1)
train_iter = I.SerialIterator(dataset_train, batchsize)
val_iter = I.SerialIterator(dataset_test,
batchsize,
repeat=False,
shuffle=False)
def scaled_abs_error(x0, x1):
if isinstance(x0, Variable):
x0 = cuda.to_cpu(x0.data)
if isinstance(x1, Variable):
x1 = cuda.to_cpu(x1.data)
scaled_x0 = ss.inverse_transform(cuda.to_cpu(x0))
scaled_x1 = ss.inverse_transform(cuda.to_cpu(x1))
diff = scaled_x0 - scaled_x1
return np.mean(np.absolute(diff), axis=0)[0]
regressor = Regressor(model,
lossfun=F.mean_squared_error,
metrics_fun={'abs_error': scaled_abs_error},
device=gpu)
optimizer = O.Adam(alpha=0.0005)
optimizer.setup(regressor)
updater = training.StandardUpdater(train_iter,
optimizer,
device=gpu,
converter=concat_mols)
dir_path = get_dir_path(batchsize, n_unit, conv_layers, M, method)
dir_path = os.path.join(dir_path, str(split_idx) + "-" + str(n_split))
os.makedirs(dir_path, exist_ok=True)
print('creating ', dir_path)
np.save(os.path.join(dir_path, "test_idx"), np.array(test_order))
trainer = training.Trainer(updater, (epoch, 'epoch'), out=dir_path)
trainer.extend(
E.Evaluator(val_iter, regressor, device=gpu, converter=concat_mols))
trainer.extend(E.LogReport())
trainer.extend(
E.PrintReport([
'epoch', 'main/loss', 'main/abs_error', 'validation/main/loss',
'validation/main/abs_error', 'elapsed_time'
]))
trainer.extend(E.ProgressBar())
trainer.run()
# --- Plot regression evaluation result ---
dataset_test = SubDataset(dataset, n_train, n, new_order)
batch_all = concat_mols(dataset_test, device=gpu)
serializers.save_npz(os.path.join(dir_path, "model.npz"), model)
result = model(batch_all[0], batch_all[1])
result = ss.inverse_transform(cuda.to_cpu(result.data))
answer = ss.inverse_transform(cuda.to_cpu(batch_all[2]))
plot_result(result,
answer,
save_filepath=os.path.join(dir_path, "result.png"))
# --- Plot regression evaluation result end ---
np.save(os.path.join(dir_path, "output.npy"), result)
np.save(os.path.join(dir_path, "answer.npy"), answer)
smiles_part = np.array(smiles)[test_order]
np.save(os.path.join(dir_path, "smiles.npy"), smiles_part)
# calculate saliency and save it.
save_result(dataset, model, dir_path, M)
def main():
# Supported preprocessing/network list
method_list = ['nfp', 'ggnn', 'schnet', 'weavenet', 'rsgcn']
scale_list = ['standardize', 'none']
parser = argparse.ArgumentParser(
description='Regression with own dataset.')
parser.add_argument('datafile', type=str)
parser.add_argument('--method', '-m', type=str, choices=method_list,
default='nfp')
parser.add_argument('--label', '-l', nargs='+',
help='target label for regression')
parser.add_argument('--scale', type=str, choices=scale_list,
default='standardize', help='Label scaling method')
parser.add_argument('--conv-layers', '-c', type=int, default=4)
parser.add_argument('--batchsize', '-b', type=int, default=32)
parser.add_argument('--gpu', '-g', type=int, default=-1)
parser.add_argument('--out', '-o', type=str, default='result')
parser.add_argument('--epoch', '-e', type=int, default=20)
parser.add_argument('--unit-num', '-u', type=int, default=16)
parser.add_argument('--seed', '-s', type=int, default=777)
parser.add_argument('--train-data-ratio', '-t', type=float, default=0.7)
args = parser.parse_args()
seed = args.seed
train_data_ratio = args.train_data_ratio
method = args.method
if args.label:
labels = args.label
class_num = len(labels) if isinstance(labels, list) else 1
else:
sys.exit("Error: No target label is specified.")
# Dataset preparation
# Postprocess is required for regression task
def postprocess_label(label_list):
return numpy.asarray(label_list, dtype=numpy.float32)
print('Preprocessing dataset...')
preprocessor = preprocess_method_dict[method]()
parser = CSVFileParser(preprocessor,
postprocess_label=postprocess_label,
labels=labels, smiles_col='SMILES')
dataset = parser.parse(args.datafile)["dataset"]
if args.scale == 'standardize':
# Standard Scaler for labels
ss = StandardScaler()
labels = ss.fit_transform(dataset.get_datasets()[-1])
dataset = NumpyTupleDataset(*(dataset.get_datasets()[:-1] + (labels,)))
train_data_size = int(len(dataset) * train_data_ratio)
train, val = split_dataset_random(dataset, train_data_size, seed)
# Network
n_unit = args.unit_num
conv_layers = args.conv_layers
if method == 'nfp':
print('Train NFP model...')
model = GraphConvPredictor(NFP(out_dim=n_unit, hidden_dim=n_unit,
n_layers=conv_layers),
MLP(out_dim=class_num, hidden_dim=n_unit))
elif method == 'ggnn':
print('Train GGNN model...')
model = GraphConvPredictor(GGNN(out_dim=n_unit, hidden_dim=n_unit,
n_layers=conv_layers),
MLP(out_dim=class_num, hidden_dim=n_unit))
elif method == 'schnet':
print('Train SchNet model...')
model = GraphConvPredictor(
SchNet(out_dim=class_num, hidden_dim=n_unit, n_layers=conv_layers),
None)
elif method == 'weavenet':
print('Train WeaveNet model...')
n_atom = 20
n_sub_layer = 1
weave_channels = [50] * conv_layers
model = GraphConvPredictor(
WeaveNet(weave_channels=weave_channels, hidden_dim=n_unit,
n_sub_layer=n_sub_layer, n_atom=n_atom),
MLP(out_dim=class_num, hidden_dim=n_unit))
elif method == 'rsgcn':
print('Train RSGCN model...')
model = GraphConvPredictor(
RSGCN(out_dim=n_unit, hidden_dim=n_unit, n_layers=conv_layers),
MLP(out_dim=class_num, hidden_dim=n_unit))
else:
raise ValueError('[ERROR] Invalid method {}'.format(method))
train_iter = I.SerialIterator(train, args.batchsize)
val_iter = I.SerialIterator(val, args.batchsize,
repeat=False, shuffle=False)
def scaled_abs_error(x0, x1):
if isinstance(x0, Variable):
x0 = cuda.to_cpu(x0.data)
if isinstance(x1, Variable):
x1 = cuda.to_cpu(x1.data)
if args.scale == 'standardize':
scaled_x0 = ss.inverse_transform(cuda.to_cpu(x0))
scaled_x1 = ss.inverse_transform(cuda.to_cpu(x1))
diff = scaled_x0 - scaled_x1
elif args.scale == 'none':
diff = cuda.to_cpu(x0) - cuda.to_cpu(x1)
return numpy.mean(numpy.absolute(diff), axis=0)[0]
classifier = L.Classifier(model, lossfun=F.mean_squared_error,
accfun=scaled_abs_error)
if args.gpu >= 0:
chainer.cuda.get_device_from_id(args.gpu).use()
classifier.to_gpu()
optimizer = O.Adam()
optimizer.setup(classifier)
updater = training.StandardUpdater(train_iter, optimizer, device=args.gpu,
converter=concat_mols)
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)
trainer.extend(E.Evaluator(val_iter, classifier, device=args.gpu,
converter=concat_mols))
trainer.extend(E.snapshot(), trigger=(args.epoch, 'epoch'))
trainer.extend(E.LogReport())
# Note that scaled errors are reported as (validation/)main/accuracy
trainer.extend(E.PrintReport(['epoch', 'main/loss', 'main/accuracy',
'validation/main/loss',
'validation/main/accuracy',
'elapsed_time']))
trainer.extend(E.ProgressBar())
trainer.run()
# Example of prediction using trained model
smiles = 'c1ccccc1'
mol = Chem.MolFromSmiles(smiles)
preprocessor = preprocess_method_dict[method]()
standardized_smiles, mol = preprocessor.prepare_smiles_and_mol(mol)
input_features = preprocessor.get_input_features(mol)
atoms, adjs = concat_mols([input_features], device=args.gpu)
prediction = model(atoms, adjs).data[0]
print('Prediction for {}:'.format(smiles))
for i, label in enumerate(args.label):
print('{}: {}'.format(label, prediction[i]))
def main():
# Parse the arguments.
args = parse_arguments()
if args.label:
labels = args.label
class_num = len(labels) if isinstance(labels, list) else 1
else:
raise ValueError('No target label was specified.')
# Dataset preparation. Postprocessing is required for the regression task.
def postprocess_label(label_list):
return numpy.asarray(label_list, dtype=numpy.float32)
# Apply a preprocessor to the dataset.
print('Preprocessing dataset...')
preprocessor = preprocess_method_dict[args.method]()
parser = CSVFileParser(preprocessor, postprocess_label=postprocess_label,
labels=labels, smiles_col='SMILES')
dataset = parser.parse(args.datafile)['dataset']
# Scale the label values, if necessary.
if args.scale == 'standardize':
scaler = StandardScaler()
labels = scaler.fit_transform(dataset.get_datasets()[-1])
dataset = NumpyTupleDataset(*(dataset.get_datasets()[:-1] + (labels,)))
else:
scaler = None
# Split the dataset into training and validation.
train_data_size = int(len(dataset) * args.train_data_ratio)
train, _ = split_dataset_random(dataset, train_data_size, args.seed)
# Set up the predictor.
predictor = set_up_predictor(args.method, args.unit_num,
args.conv_layers, class_num)
# Set up the iterator.
train_iter = SerialIterator(train, args.batchsize)
# Set up the regressor.
metrics_fun = {'mean_abs_error': MeanAbsError(scaler=scaler),
'root_mean_sqr_error': RootMeanSqrError(scaler=scaler)}
regressor = Regressor(predictor, lossfun=F.mean_squared_error,
metrics_fun=metrics_fun, device=args.gpu)
# Set up the optimizer.
optimizer = optimizers.Adam()
optimizer.setup(regressor)
# Set up the updater.
updater = training.StandardUpdater(train_iter, optimizer, device=args.gpu,
converter=concat_mols)
# Set up the trainer.
print('Training...')
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)
trainer.extend(E.snapshot(), trigger=(args.epoch, 'epoch'))
trainer.extend(E.LogReport())
trainer.extend(E.PrintReport(['epoch', 'main/loss', 'main/mean_abs_error',
'main/root_mean_sqr_error', 'elapsed_time']))
trainer.extend(E.ProgressBar())
trainer.run()
# Save the regressor's parameters.
model_path = os.path.join(args.out, args.model_filename)
print('Saving the trained model to {}...'.format(model_path))
regressor.save_pickle(model_path, protocol=args.protocol)
# Save the standard scaler's parameters.
if scaler is not None:
with open(os.path.join(args.out, 'scaler.pkl'), mode='wb') as f:
pickle.dump(scaler, f, protocol=args.protocol)
def main():
# Parse the arguments.
args = parse_arguments()
# Set up some useful variables that will be used later on.
method = args.method
if args.label != 'all':
labels = args.label
cache_dir = os.path.join('input', '{}_{}'.format(method, labels))
class_num = len(labels) if isinstance(labels, list) else 1
else:
labels = None
cache_dir = os.path.join('input', '{}_all'.format(method))
class_num = len(D.get_qm9_label_names())
# Get the filename corresponding to the cached dataset, based on the amount
# of data samples that need to be parsed from the original dataset.
num_data = args.num_data
if num_data >= 0:
dataset_filename = 'data_{}.npz'.format(num_data)
else:
dataset_filename = 'data.npz'
# Load the cached dataset.
dataset_cache_path = os.path.join(cache_dir, dataset_filename)
dataset = None
if os.path.exists(dataset_cache_path):
print('Loading cached dataset from {}.'.format(dataset_cache_path))
dataset = NumpyTupleDataset.load(dataset_cache_path)
if dataset is None:
print('Preprocessing dataset...')
preprocessor = preprocess_method_dict[method]()
if num_data >= 0:
# Select the first `num_data` samples from the dataset.
target_index = numpy.arange(num_data)
dataset = D.get_qm9(preprocessor, labels=labels,
target_index=target_index)
else:
# Load the entire dataset.
dataset = D.get_qm9(preprocessor, labels=labels)
# Cache the laded dataset.
if not os.path.exists(cache_dir):
os.makedirs(cache_dir)
NumpyTupleDataset.save(dataset_cache_path, dataset)
# Scale the label values, if necessary.
if args.scale == 'standardize':
print('Applying standard scaling to the labels.')
scaler = StandardScaler()
scaled_t = scaler.fit_transform(dataset.get_datasets()[-1])
dataset = NumpyTupleDataset(*(dataset.get_datasets()[:-1]
+ (scaled_t,)))
else:
print('No standard scaling was selected.')
scaler = None
# Split the dataset into training and validation.
train_data_size = int(len(dataset) * args.train_data_ratio)
train, valid = split_dataset_random(dataset, train_data_size, args.seed)
# Set up the predictor.
predictor = set_up_predictor(method, args.unit_num, args.conv_layers,
class_num, scaler)
# Set up the iterators.
train_iter = iterators.SerialIterator(train, args.batchsize)
valid_iter = iterators.SerialIterator(valid, args.batchsize, repeat=False,
shuffle=False)
# Set up the regressor.
device = args.gpu
metrics_fun = {'mae': MeanAbsError(scaler=scaler),
'rmse': RootMeanSqrError(scaler=scaler)}
regressor = Regressor(predictor, lossfun=F.mean_squared_error,
metrics_fun=metrics_fun, device=device)
# Set up the optimizer.
optimizer = optimizers.Adam()
optimizer.setup(regressor)
# Set up the updater.
updater = training.StandardUpdater(train_iter, optimizer, device=device,
converter=concat_mols)
# Set up the trainer.
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)
trainer.extend(E.Evaluator(valid_iter, regressor, device=device,
converter=concat_mols))
trainer.extend(E.snapshot(), trigger=(args.epoch, 'epoch'))
trainer.extend(E.LogReport())
trainer.extend(E.PrintReport([
'epoch', 'main/loss', 'main/mae', 'main/rmse', 'validation/main/loss',
'validation/main/mae', 'validation/main/rmse', 'elapsed_time']))
trainer.extend(E.ProgressBar())
trainer.run()
# Save the regressor's parameters.
model_path = os.path.join(args.out, args.model_filename)
print('Saving the trained model to {}...'.format(model_path))
regressor.save_pickle(model_path, protocol=args.protocol)