def build_predictor(method, n_unit, conv_layers, class_num):
if method == 'nfp':
print('Use NFP predictor...')
predictor = 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('Use GGNN predictor...')
predictor = 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('Use SchNet predictor...')
predictor = SchNet(out_dim=class_num,
hidden_dim=n_unit,
n_layers=conv_layers,
readout_hidden_dim=n_unit)
elif method == 'weavenet':
print('Use WeaveNet predictor...')
n_atom = 20
n_sub_layer = 1
weave_channels = [50] * conv_layers
predictor = 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))
else:
raise ValueError('[ERROR] Invalid predictor: method={}'.format(method))
return predictor
def set_up_predictor(method, n_unit, conv_layers, class_num, scaler):
"""Sets up the predictor, consisting of a graph convolution network and
a multilayer perceptron.
Args:
method (str): Method name.
n_unit (int): Number of hidden units.
conv_layers (int): Number of convolutional layers for the graph
convolution network.
class_num (int): Number of output classes.
Returns:
predictor (chainer.Chain): An instance of the selected predictor.
"""
mlp = MLP(out_dim=class_num, hidden_dim=n_unit)
if method == 'nfp':
print('Training an NFP predictor...')
nfp = NFP(out_dim=n_unit, hidden_dim=n_unit, n_layers=conv_layers)
predictor = GraphConvPredictor(nfp, mlp, scaler)
elif method == 'ggnn':
print('Training a GGNN predictor...')
ggnn = GGNN(out_dim=n_unit, hidden_dim=n_unit, n_layers=conv_layers)
predictor = GraphConvPredictor(ggnn, mlp, scaler)
elif method == 'schnet':
print('Training an SchNet predictor...')
schnet = SchNet(out_dim=class_num,
hidden_dim=n_unit,
n_layers=conv_layers)
predictor = GraphConvPredictor(schnet, None, scaler)
elif method == 'weavenet':
print('Training a WeaveNet predictor...')
n_atom = 20
n_sub_layer = 1
weave_channels = [50] * conv_layers
weavenet = WeaveNet(weave_channels=weave_channels,
hidden_dim=n_unit,
n_sub_layer=n_sub_layer,
n_atom=n_atom)
predictor = GraphConvPredictor(weavenet, mlp, scaler)
elif method == 'rsgcn':
print('Training an RSGCN predictor...')
rsgcn = RSGCN(out_dim=n_unit, hidden_dim=n_unit, n_layers=conv_layers)
predictor = GraphConvPredictor(rsgcn, mlp, scaler)
elif method == 'relgcn':
print('Use Relational GCN predictor...')
num_edge_type = 4
relgcn = RelGCN(out_channels=n_unit,
num_edge_type=num_edge_type,
scale_adj=True)
predictor = GraphConvPredictor(relgcn, mlp, scaler)
elif method == 'relgat':
print('Train Relational GAT predictor...')
relgat = RelGAT(out_dim=n_unit,
hidden_dim=n_unit,
n_layers=conv_layers)
predictor = GraphConvPredictor(relgat, mlp, scaler)
else:
raise ValueError('[ERROR] Invalid method: {}'.format(method))
return predictor
def set_up_predictor(method,
fp_hidden_dim,
fp_out_dim,
conv_layers,
concat_hidden,
fp_dropout_rate,
net_hidden_dims,
class_num,
sim_method='mlp'):
if sim_method == 'mlp':
mlp = MLP(out_dim=class_num, hidden_dims=net_hidden_dims)
elif sim_method == 'cosine':
mlp = Cosine()
else:
raise ValueError(
'[ERROR] Invalid similarity method: {}'.format(method))
print('Using {} as similarity predictor with hidden_dims {}...'.format(
sim_method, net_hidden_dims))
if method == 'nfp':
print('Training an NFP predictor...')
nfp = NFP(out_dim=fp_out_dim,
hidden_dim=fp_hidden_dim,
n_layers=conv_layers,
concat_hidden=concat_hidden)
predictor = GraphConvPredictorForPair(nfp, mlp)
elif method == 'ggnn':
print('Training a GGNN predictor...')
ggnn = GGNN(out_dim=fp_out_dim,
hidden_dim=fp_hidden_dim,
n_layers=conv_layers,
concat_hidden=concat_hidden,
dropout_rate=fp_dropout_rate)
predictor = GraphConvPredictorForPair(ggnn, mlp)
elif method == 'schnet':
print('Training an SchNet predictor...')
schnet = SchNet(out_dim=class_num,
hidden_dim=fp_hidden_dim,
n_layers=conv_layers)
predictor = GraphConvPredictorForPair(schnet, None)
elif method == 'weavenet':
print('Training a WeaveNet predictor...')
n_atom = 20
n_sub_layer = 1
weave_channels = [50] * conv_layers
weavenet = WeaveNet(weave_channels=weave_channels,
hidden_dim=fp_hidden_dim,
n_sub_layer=n_sub_layer,
n_atom=n_atom)
predictor = GraphConvPredictorForPair(weavenet, mlp)
elif method == 'rsgcn':
print('Training an RSGCN predictor...')
rsgcn = RSGCN(out_dim=fp_out_dim,
hidden_dim=fp_hidden_dim,
n_layers=conv_layers)
predictor = GraphConvPredictorForPair(rsgcn, mlp)
else:
raise ValueError('[ERROR] Invalid method: {}'.format(method))
return predictor
def set_up_predictor(method, n_unit, conv_layers, class_num):
"""Sets up the graph convolution network predictor.
Args:
method: Method name. Currently, the supported ones are `nfp`, `ggnn`,
`schnet`, `weavenet` and `rsgcn`.
n_unit: Number of hidden units.
conv_layers: Number of convolutional layers for the graph convolution
network.
class_num: Number of output classes.
Returns:
An instance of the selected predictor.
"""
predictor = None
mlp = MLP(out_dim=class_num, hidden_dim=n_unit)
if method == 'nfp':
print('Training an NFP predictor...')
nfp = NFP(out_dim=n_unit, hidden_dim=n_unit, n_layers=conv_layers)
predictor = GraphConvPredictor(nfp, mlp)
elif method == 'ggnn':
print('Training a GGNN predictor...')
ggnn = GGNN(out_dim=n_unit, hidden_dim=n_unit, n_layers=conv_layers)
predictor = GraphConvPredictor(ggnn, mlp)
elif method == 'schnet':
print('Training an SchNet predictor...')
schnet = SchNet(out_dim=class_num,
hidden_dim=n_unit,
n_layers=conv_layers)
predictor = GraphConvPredictor(schnet, None)
elif method == 'weavenet':
print('Training a WeaveNet predictor...')
n_atom = 20
n_sub_layer = 1
weave_channels = [50] * conv_layers
weavenet = WeaveNet(weave_channels=weave_channels,
hidden_dim=n_unit,
n_sub_layer=n_sub_layer,
n_atom=n_atom)
predictor = GraphConvPredictor(weavenet, mlp)
elif method == 'rsgcn':
print('Training an RSGCN predictor...')
rsgcn = RSGCN(out_dim=n_unit, hidden_dim=n_unit, n_layers=conv_layers)
predictor = GraphConvPredictor(rsgcn, mlp)
elif method == 'relgcn':
print('Training an RelGCN predictor...')
num_edge_type = 4
relgcn = RelGCN(out_channels=class_num,
num_edge_type=num_edge_type,
scale_adj=True)
predictor = GraphConvPredictor(relgcn, None)
else:
raise ValueError('[ERROR] Invalid method: {}'.format(method))
return predictor
def build_predictor(method, n_unit, conv_layers, class_num):
if method == 'nfp':
print('Use NFP predictor...')
predictor = 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('Use GGNN predictor...')
predictor = 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('Use SchNet predictor...')
# MLP layer is not necessary for SchNet
predictor = GraphConvPredictor(
SchNet(out_dim=class_num,
hidden_dim=n_unit,
n_layers=conv_layers,
readout_hidden_dim=n_unit), None)
elif method == 'weavenet':
print('Use WeaveNet predictor...')
n_atom = 20
n_sub_layer = 1
weave_channels = [50] * conv_layers
predictor = 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('Use RSGCN predictor...')
predictor = GraphConvPredictor(
RSGCN(out_dim=n_unit, hidden_dim=n_unit, n_layers=conv_layers),
MLP(out_dim=class_num, hidden_dim=n_unit))
elif method == 'relgcn':
print('Use Relational GCN predictor...')
num_edge_type = 4
predictor = GraphConvPredictor(
RelGCN(out_channels=n_unit,
num_edge_type=num_edge_type,
scale_adj=True), MLP(out_dim=class_num, hidden_dim=n_unit))
elif method == 'relgat':
print('Use GAT predictor...')
predictor = GraphConvPredictor(
RelGAT(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 predictor: method={}'.format(method))
return predictor
def main():
"""Launcher."""
preprocessor = preprocess_method_dict["nfp"]()
dataset = datasets.get_qm9(preprocessor, labels="h**o")
cache_dir = "data/"
if not (os.path.exists(cache_dir)):
os.makedirs(cache_dir)
NumpyTupleDataset.save(cache_dir + "data.npz", dataset)
dataset = NumpyTupleDataset.load(cache_dir + 'data.npz')
train_data_ratio = 0.7
train_data_size = int(len(dataset) * train_data_ratio)
train, validation = split_dataset_random(dataset, train_data_size, 777)
print('train dataset size:', len(train))
print('validation dataset size:', len(validation))
n_unit = 16
conv_layers = 4
model = GraphConvPredictor(NFP(n_unit, n_unit, conv_layers),
MLP(n_unit, 1))
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 set_up_predictor(method, n_unit, conv_layers, class_num):
"""Sets up the predictor, consisting of a graph convolution network and
a multilayer perceptron.
Args:
method: Method name. See `parse_arguments`.
n_unit: Number of hidden units.
conv_layers: Number of convolutional layers for the graph convolution
network.
class_num: Number of output classes.
Returns:
An instance of the selected predictor.
"""
mlp = MLP(out_dim=class_num, hidden_dim=n_unit)
if method == 'nfp':
print('Training an NFP predictor...')
nfp = NFP(out_dim=n_unit, hidden_dim=n_unit, n_layers=conv_layers)
return GraphConvPredictor(nfp, mlp)
elif method == 'nfp_gwm':
print('Training an NFP+GWM predictor...')
nfp_gwm = NFP_GWM(out_dim=n_unit,
hidden_dim=n_unit,
hidden_dim_super=n_unit,
n_layers=conv_layers,
dropout_ratio=0.5)
return GraphConvPredictorForGWM(nfp_gwm, mlp)
elif method == 'ggnn':
print('Training a GGNN predictor...')
ggnn = GGNN(out_dim=n_unit, hidden_dim=n_unit, n_layers=conv_layers)
return GraphConvPredictor(ggnn, mlp)
elif method == 'ggnn_gwm':
print('Train GGNN+GWM model...')
ggnn_gwm = GGNN_GWM(out_dim=n_unit,
hidden_dim=n_unit,
hidden_dim_super=n_unit,
n_layers=conv_layers,
dropout_ratio=0.5,
weight_tying=True)
return GraphConvPredictorForGWM(ggnn_gwm, mlp)
elif method == 'schnet':
print('Training an SchNet predictor...')
schnet = SchNet(out_dim=class_num,
hidden_dim=n_unit,
n_layers=conv_layers)
return GraphConvPredictor(schnet, None)
elif method == 'weavenet':
print('Training a WeaveNet predictor...')
n_atom = 20
n_sub_layer = 1
weave_channels = [50] * conv_layers
weavenet = WeaveNet(weave_channels=weave_channels,
hidden_dim=n_unit,
n_sub_layer=n_sub_layer,
n_atom=n_atom)
return GraphConvPredictor(weavenet, mlp)
elif method == 'rsgcn':
print('Training an RSGCN predictor...')
rsgcn = RSGCN(out_dim=n_unit, hidden_dim=n_unit, n_layers=conv_layers)
return GraphConvPredictor(rsgcn, mlp)
elif method == 'rsgcn_gwm':
print('Training an RSGCN+GWM predictor...')
rsgcn_gwm = RSGCN_GWM(out_dim=n_unit,
hidden_dim=n_unit,
hidden_dim_super=n_unit,
n_layers=conv_layers,
dropout_ratio=0.5)
return GraphConvPredictorForGWM(rsgcn_gwm, mlp)
elif method == 'relgcn':
print('Training an RelGCN predictor...')
num_edge_type = 4
relgcn = RelGCN(out_channels=n_unit,
num_edge_type=num_edge_type,
scale_adj=True)
return GraphConvPredictor(relgcn, mlp)
elif method == 'relgat':
print('Train Relational GAT model...')
relgat = RelGAT(out_dim=n_unit,
hidden_dim=n_unit,
n_layers=conv_layers)
return GraphConvPredictor(relgat, mlp)
elif method == 'gin':
print('Training a GIN predictor...')
gin = GIN(out_dim=n_unit, hidden_dim=n_unit, n_layers=conv_layers)
return GraphConvPredictor(gin, mlp)
elif method == 'gin_gwm':
print('Training a GIN+GWM predictor...')
gin_gwm = GIN_GWM(out_dim=n_unit,
hidden_dim=n_unit,
hidden_dim_super=n_unit,
n_layers=conv_layers,
dropout_ratio=0.5,
weight_tying=True)
return GraphConvPredictorForGWM(gin_gwm, mlp)
raise ValueError('[ERROR] Invalid method: {}'.format(method))
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']
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)
val_sampler = SubsetRandomSampler(val_indices)
loader = DataLoader(generator_train,
batch_size=BATCH_SIZE,
sampler=train_sampler)
loader_val = DataLoader(generator_train,
batch_size=BATCH_SIZE,
sampler=val_sampler)
loader_test = DataLoader(generator_test, batch_size=BATCH_SIZE)
#create models
if model_name == "baseline":
model = BaselineModel(input_len, drop_rate=dropout)
elif model_name == 'nfp':
print('Training an NFP predictor...')
mlp = MLP(out_dim=class_num, hidden_dim=n_unit)
nfp = NFP(out_dim=n_unit, hidden_dim=n_unit, n_layers=conv_layers)
model = GraphConvPredictor(nfp, mlp)
elif model_name == 'ggnn':
print('Training a GGNN predictor...')
mlp = MLP(out_dim=class_num, hidden_dim=n_unit)
ggnn = GGNN(out_dim=n_unit, hidden_dim=n_unit, n_layers=conv_layers)
model = GraphConvPredictor(ggnn, mlp)
elif model_name == 'schnet':
print('Training an SchNet predictor...')
mlp = MLP(out_dim=class_num, hidden_dim=n_unit)
schnet = SchNet(out_dim=class_num, hidden_dim=n_unit, n_layers=conv_layers)
model = GraphConvPredictor(schnet, None)
elif model_name == 'weavenet':
print('Training a WeaveNet predictor...')
mlp = MLP(out_dim=class_num, hidden_dim=n_unit)
n_atom = 20
def set_up_predictor(method, fp_hidden_dim, fp_out_dim, conv_layers, concat_hidden, layer_aggregator,
fp_dropout_rate, fp_batch_normalization,
net_hidden_dims, class_num,
sim_method='mlp', fp_attention=False, weight_typing=True, attention_tying=True,
update_attention=False,
context=False, context_layers=1, context_dropout=0.,
message_function='matrix_multiply', readout_function='graph_level',
num_timesteps=3, num_output_hidden_layers=0, output_hidden_dim=16,
output_activation=functions.relu,
symmetric=None,
):
if sim_method == 'mlp':
logging.info('Using multi-layer perceptron for the learning of composite representation')
mlp = MLP(out_dim=class_num, hidden_dims=net_hidden_dims)
elif sim_method == 'ntn':
logging.info('Using neural tensor network for the learning of composite representation')
ntn_out_dim = 8
logging.info('NTN out dim is {}'.format(ntn_out_dim))
mlp = NTN(left_dim=fp_out_dim, right_dim=fp_out_dim, out_dim=class_num,
ntn_out_dim=ntn_out_dim, hidden_dims=net_hidden_dims)
elif sim_method == 'symmlp':
logging.info('Using symmetric multi-layer perceptron for the learning of composite representation')
mlp = MLP(out_dim=class_num, hidden_dims=net_hidden_dims)
elif sim_method == 'hole':
logging.info('Using holegraphic embedding for the learning of composite representation')
mlp = HolE(out_dim=class_num, hidden_dims=net_hidden_dims)
elif sim_method == 'dist-mult':
logging.info('Using DistMult embedding for the learning of composite representation')
dm_out_dim = 8
mlp = DistMult(left_dim=fp_out_dim, right_dim=fp_out_dim, out_dim=class_num,
dm_out_dim=dm_out_dim, hidden_dims=net_hidden_dims)
else:
raise ValueError('[ERROR] Invalid similarity method: {}'.format(method))
logging.info('Using {} as similarity predictor with hidden_dims {}...'.format(sim_method, net_hidden_dims))
encoder = None
if method == 'ggnn':
logging.info('Training a GGNN predictor...')
if fp_attention:
logging.info('Self-attention mechanism is utilized...')
if attention_tying:
logging.info('Self-attention is tying...')
else:
logging.info('Self-attention is not tying...')
if update_attention:
logging.info('Self-attention mechanism is utilized in update...')
if attention_tying:
logging.info('Self-attention is tying...')
else:
logging.info('Self-attention is not tying...')
if not weight_typing:
logging.info('Weight is not tying...')
if fp_dropout_rate != 0.0:
logging.info('Using dropout whose rate is {}...'.format(fp_dropout_rate))
if fp_batch_normalization:
logging.info('Using batch normalization in dynamic fingerprint...')
if concat_hidden:
logging.info('Incorporating layer aggregation via concatenation after readout...')
if layer_aggregator:
logging.info('Incorporating layer aggregation via {} before readout...'.format(layer_aggregator))
if context:
logging.info('Context embedding is utilized...')
logging.info('Number of context layers is {}...'.format(context_layers))
logging.info('Dropout rate of context layers is {:.2f}'.format(context_dropout))
logging.info('Message function is {}'.format(message_function))
logging.info('Readout function is {}'.format(readout_function))
logging.info('Num_timesteps = {}, num_output_hidden_layers={}, output_hidden_dim={}'.format(
num_timesteps, num_output_hidden_layers, output_hidden_dim
))
# num_timesteps=3, num_output_hidden_layers=0, output_hidden_dim=16, output_activation=functions.relu,
encoder = GGNN(out_dim=fp_out_dim, hidden_dim=fp_hidden_dim, n_layers=conv_layers, concat_hidden=concat_hidden,
layer_aggregator=layer_aggregator,
dropout_rate=fp_dropout_rate, batch_normalization=fp_batch_normalization,
use_attention=fp_attention, weight_tying=weight_typing, attention_tying=attention_tying,
context=context, message_function=message_function, readout_function=readout_function,
num_timesteps=num_timesteps, num_output_hidden_layers=num_output_hidden_layers,
output_hidden_dim=output_hidden_dim, output_activation=output_activation)
elif method == 'nfp':
fp_max_degree = 6
logging.info('Training an NFP predictor...')
logging.info('Max degree is {}'.format(fp_max_degree))
encoder = NFP(out_dim=fp_out_dim, hidden_dim=fp_hidden_dim, n_layers=conv_layers, concat_hidden=concat_hidden,
max_degree=fp_max_degree)
elif method == 'schnet':
logging.info('Training an SchNet predictor...')
schnet = SchNet(out_dim=class_num, hidden_dim=fp_hidden_dim,
n_layers=conv_layers)
encoder = schnet
elif method == 'weavenet':
logging.info('Training a WeaveNet predictor...')
n_atom = 20
n_sub_layer = 1
weave_channels = [50] * conv_layers
encoder = WeaveNet(weave_channels=weave_channels, hidden_dim=fp_hidden_dim,
n_sub_layer=n_sub_layer, n_atom=n_atom)
elif method == 'rsgcn':
logging.info('Training an RSGCN predictor...')
use_batch_norm = True
dropout_ratio = 0.5
if use_batch_norm:
logging.info('Using batch normalization...')
logging.info('Dropout ratio is {:.1f}'.format(dropout_ratio))
encoder = RSGCN(out_dim=fp_out_dim, hidden_dim=fp_hidden_dim, n_layers=conv_layers, use_batch_norm=use_batch_norm, dropout_ratio=dropout_ratio)
elif method == 'gin':
encoder = GIN(out_dim=fp_out_dim, hidden_dim=fp_hidden_dim, n_layers=conv_layers,
dropout_ratio=0.5, concat_hidden=True)
else:
raise ValueError('[ERROR] Invalid method: {}'.format(method))
predictor = GraphConvPredictorForPair(encoder, mlp, symmetric=symmetric)
return predictor
def main():
method_list = ['nfp', 'ggnn', 'schnet', 'weavenet', 'rsgcn']
dataset_names = list(molnet_default_config.keys())
parser = argparse.ArgumentParser(description='molnet example')
parser.add_argument('--method',
'-m',
type=str,
choices=method_list,
default='nfp')
parser.add_argument('--label',
'-l',
type=str,
default='',
help='target label for regression, empty string means '
'to predict all property at once')
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('--dataset',
'-d',
type=str,
choices=dataset_names,
default='bbbp')
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()
dataset_name = args.dataset
method = args.method
num_data = args.num_data
n_unit = args.unit_num
conv_layers = args.conv_layers
print('Use {} dataset'.format(dataset_name))
if args.label:
labels = args.label
cache_dir = os.path.join(
'input', '{}_{}_{}'.format(dataset_name, method, labels))
class_num = len(labels) if isinstance(labels, list) else 1
else:
labels = None
cache_dir = os.path.join('input',
'{}_{}_all'.format(dataset_name, method))
class_num = len(molnet_default_config[args.dataset]['tasks'])
# Dataset preparation
def get_dataset_paths(cache_dir, num_data):
filepaths = []
for filetype in ['train', 'valid', 'test']:
filename = filetype + '_data'
if num_data >= 0:
filename += '_' + str(num_data)
filename += '.npz'
filepath = os.path.join(cache_dir, filename)
filepaths.append(filepath)
return filepaths
filepaths = get_dataset_paths(cache_dir, num_data)
if all([os.path.exists(fpath) for fpath in filepaths]):
datasets = []
for fpath in filepaths:
print('load from cache {}'.format(fpath))
datasets.append(NumpyTupleDataset.load(fpath))
else:
print('preprocessing dataset...')
preprocessor = preprocess_method_dict[method]()
# only use first 100 for debug if num_data >= 0
target_index = numpy.arange(num_data) if num_data >= 0 else None
datasets = D.molnet.get_molnet_dataset(dataset_name,
preprocessor,
labels=labels,
target_index=target_index)
if not os.path.exists(cache_dir):
os.makedirs(cache_dir)
datasets = datasets['dataset']
for i, fpath in enumerate(filepaths):
NumpyTupleDataset.save(fpath, datasets[i])
train, val, _ = datasets
# Network
if method == 'nfp':
print('Train NFP model...')
predictor = 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...')
predictor = 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...')
predictor = 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
predictor = 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...')
predictor = 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)
metrics_fun = molnet_default_config[dataset_name]['metrics']
loss_fun = molnet_default_config[dataset_name]['loss']
task_type = molnet_default_config[dataset_name]['task_type']
if task_type == 'regression':
model = Regressor(predictor,
lossfun=loss_fun,
metrics_fun=metrics_fun,
device=args.gpu)
# TODO(nakago): Use standard scaler for regression task
elif task_type == 'classification':
model = Classifier(predictor,
lossfun=loss_fun,
metrics_fun=metrics_fun,
device=args.gpu)
else:
raise NotImplementedError(
'Not implemented task_type = {}'.format(task_type))
optimizer = optimizers.Adam()
optimizer.setup(model)
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, model, device=args.gpu, converter=concat_mols))
trainer.extend(E.snapshot(), trigger=(args.epoch, 'epoch'))
trainer.extend(E.LogReport())
print_report_targets = ['epoch', 'main/loss', 'validation/main/loss']
if metrics_fun is not None and type(metrics_fun) == dict:
for m_k in metrics_fun.keys():
print_report_targets.append('main/' + m_k)
print_report_targets.append('validation/main/' + m_k)
if task_type == 'classification':
# Evaluation for train data takes time, skip for now.
# trainer.extend(ROCAUCEvaluator(
# train_iter, model, device=args.gpu, eval_func=predictor,
# converter=concat_mols, name='train', raise_value_error=False))
# print_report_targets.append('train/main/roc_auc')
trainer.extend(
ROCAUCEvaluator(val_iter,
model,
device=args.gpu,
eval_func=predictor,
converter=concat_mols,
name='val',
raise_value_error=False))
print_report_targets.append('val/main/roc_auc')
print_report_targets.append('elapsed_time')
trainer.extend(E.PrintReport(print_report_targets))
trainer.extend(E.ProgressBar())
trainer.run()
# --- save model ---
protocol = args.protocol
model.save_pickle(os.path.join(args.out, args.model_filename),
protocol=protocol)
def set_up_predictor(method,
fp_hidden_dim,
fp_out_dim,
conv_layers,
concat_hidden,
fp_dropout_rate,
fp_batch_normalization,
net_hidden_dims,
class_num,
weight_typing=True,
sim_method='mlp',
symmetric=None,
attn_model=None):
sim_method_dict = {
'mlp': 'multi-layered perceptron',
'ntn': 'bilinear transform',
'symmlp': 'symmetric perceptron',
'hole': 'holographic embedding',
'dist-mult': 'dist-mult',
}
method_dict = {
'ggnn': 'GGNN',
'nfp': 'NFP',
}
logging.info('Graph Embedding: {}'.format(method_dict.get(method, None)))
logging.info('Link Prediction: {}'.format(
sim_method_dict.get(sim_method, None)))
lp = None
if sim_method == 'mlp':
lp = MLP(out_dim=class_num, hidden_dims=net_hidden_dims)
elif sim_method == 'ntn':
ntn_out_dim = 8
lp = NTN(left_dim=fp_out_dim,
right_dim=fp_out_dim,
out_dim=class_num,
ntn_out_dim=ntn_out_dim,
hidden_dims=net_hidden_dims)
elif sim_method == 'symmlp':
lp = MLP(out_dim=class_num, hidden_dims=net_hidden_dims)
elif sim_method == 'hole':
lp = HolE(out_dim=class_num, hidden_dims=net_hidden_dims)
elif sim_method == 'dist-mult':
dm_out_dim = 8
lp = DistMult(left_dim=fp_out_dim,
right_dim=fp_out_dim,
out_dim=class_num,
dm_out_dim=dm_out_dim,
hidden_dims=net_hidden_dims)
else:
raise ValueError(
'[ERROR] Invalid link prediction model: {}'.format(method))
attn = None
scorer = 'bilinear'
if attn_model == 'alter':
attn_weight_tying = True
logging.info('Using alternating co-attention')
if attn_weight_tying:
logging.info('Weight is tying')
attn = AlternatingCoattention(hidden_dim=fp_hidden_dim,
out_dim=fp_out_dim,
head=8,
weight_tying=True)
elif attn_model == 'para':
attn_weight_tying = True
logging.info('Using parallel co-attention')
logging.info('Scorer is {}'.format(scorer))
if attn_weight_tying:
logging.info('Weight is tying')
attn = ParallelCoattention(hidden_dim=fp_hidden_dim,
out_dim=fp_out_dim,
head=1,
activation=F.tanh,
weight_tying=attn_weight_tying)
elif attn_model == 'circ':
logging.info('Using circular based parallel co-attention')
attn = CircularParallelCoattention(hidden_dim=fp_hidden_dim,
out_dim=fp_out_dim,
activation=F.tanh)
elif attn_model == 'vqa':
logging.info('Using vqa fine-grained co-attention')
attn = VQAParallelCoattention(hidden_dim=fp_hidden_dim * 2,
out_dim=fp_out_dim,
head=8)
elif attn_model == 'pool':
logging.info('Using pool fine-graind co-attention')
attn = PoolingFineCoattention(hidden_dim=fp_hidden_dim * 2,
out_dim=fp_out_dim)
elif attn_model == 'lt':
logging.info('Using lt fine-grained co-attention')
attn = LinearTransformFineCoattention(hidden_dim=fp_hidden_dim * 2,
out_dim=fp_out_dim)
elif attn_model == 'nie':
logging.info('Using nie fine-grained co-attention')
logging.info('Using activation function tanh')
# multiplied by 2 for concat
attn = NieFineCoattention(hidden_dim=fp_hidden_dim * 2,
out_dim=fp_out_dim,
head=8,
activation=F.tanh)
elif attn_model == 'bimpm':
logging.info('Using bimpm matching strategy')
attn = BiMPM(hidden_dim=fp_hidden_dim,
out_dim=fp_out_dim,
head=fp_out_dim,
with_max_pool=True,
with_att_mean=True,
with_att_max=True,
aggr=F.sum)
elif attn_model == 'global':
logging.info('Using global coattention')
attn = GlobalCoattention(hidden_dim=fp_hidden_dim,
out_dim=fp_out_dim,
weight_tying=True)
elif attn_model == 'neural':
logging.info('Using neural coattention')
attn = NeuralCoattention(hidden_dim=fp_hidden_dim,
out_dim=fp_out_dim,
weight_tying=True)
else:
raise ValueError('[ERROR] Invalid Co-Attention Method.')
siamese = True
if not siamese:
encoder_1, encoder_2 = None, None
else:
encoder = None
if method == 'ggnn':
if not weight_typing:
logging.info('Weight is not tying')
if fp_dropout_rate != 0.0:
logging.info('Forward propagation dropout rate is {:.1f}'.format(
fp_dropout_rate))
if fp_batch_normalization:
logging.info('Using batch normalization')
if concat_hidden:
logging.info('Using concatenation between layers')
if not siamese:
encoder_1 = GGNN(out_dim=fp_out_dim,
hidden_dim=fp_hidden_dim,
n_layers=conv_layers,
concat_hidden=concat_hidden,
weight_tying=weight_typing)
encoder_2 = GGNN(out_dim=fp_out_dim,
hidden_dim=fp_hidden_dim,
n_layers=conv_layers,
concat_hidden=concat_hidden,
weight_tying=weight_typing)
else:
encoder = GGNN(out_dim=fp_out_dim,
hidden_dim=fp_hidden_dim,
n_layers=conv_layers,
concat_hidden=concat_hidden,
weight_tying=weight_typing)
elif method == 'nfp':
print('Training an NFP predictor...')
encoder = NFP(out_dim=fp_out_dim,
hidden_dim=fp_hidden_dim,
n_layers=conv_layers,
concat_hidden=concat_hidden)
else:
raise ValueError('[ERROR] Invalid graph embedding encoder.')
if siamese:
predictor = GraphConvPredictorForPair(encoder,
attn,
lp,
symmetric=symmetric,
siamese=siamese,
use_s_lstm=True,
use_i_lstm=True)
else:
predictor = GraphConvPredictorForPair(encoder_1,
attn,
lp,
symmetric=symmetric,
siamese=siamese,
another_graph_conv=encoder_2)
return predictor
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]))