def test_save_load(self, data): tmp_cache_path = os.path.join(tempfile.mkdtemp(), 'tmp.npz') dataset = NumpyTupleDataset(*data) NumpyTupleDataset.save(tmp_cache_path, dataset) assert os.path.exists(tmp_cache_path) load_dataset = NumpyTupleDataset.load(tmp_cache_path) os.remove(tmp_cache_path) assert len(dataset._datasets) == len(load_dataset._datasets) for a, d in six.moves.zip(dataset._datasets, load_dataset._datasets): numpy.testing.assert_array_equal(a, d)
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'] 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()
if __name__ == "__main__": # -- get dataset -- # parser = argparse.ArgumentParser() parser.add_argument("-p", "--path", type=str, default=None, help="path to your dataset") parser.add_argument("-a", "--atomic_num_path", type=str, default=atomic_num_list_path, help="path to your atomic num list file") args = parser.parse_args() dataset = NumpyTupleDataset.load(args.path) atomic_num_list = load_id_to_atomic_num(args.atomic_num_path) periodic_table = load_periodic_table() result_dict = {} for x, _ in tqdm(dataset): atomic_ids = filter(lambda x: x > 0, map(lambda x: atomic_num_list[x], x.astype(int))) for a in atomic_ids: symbol = periodic_table.loc[a]["symbol"] result_dict[symbol] = result_dict[ symbol] + 1 if symbol in result_dict else 1 for k, v in result_dict.items(): print("{:<5}{:<5}".format(k, v))
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: labels = args.label cache_dir = os.path.join('input', '{}_{}'.format(method, labels)) 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) # Load the standard scaler parameters, if necessary. if args.scale == 'standardize': scaler_path = os.path.join(args.in_dir, 'scaler.pkl') print('Loading scaler parameters from {}.'.format(scaler_path)) with open(scaler_path, mode='rb') as f: scaler = pickle.load(f) else: print('No standard scaling was selected.') scaler = None # 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) # 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) # Replace the default predictor with one that scales the output labels. scaled_predictor = ScaledGraphConvPredictor(regressor.predictor) scaled_predictor.scaler = scaler regressor.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. print('Predicting...') y_pred = regressor.predict(test, converter=extract_inputs) # Extract the ground-truth labels. t = concat_mols(test, device=-1)[-1] n_eval = 10 # Construct dataframe. df_dict = {} for i, l in enumerate(labels): df_dict.update({ 'y_pred_{}'.format(l): y_pred[:, i], 't_{}'.format(l): t[:, i], }) df = pandas.DataFrame(df_dict) # Show a prediction/ground truth table with 5 random examples. print(df.sample(5)) for target_label in range(y_pred.shape[1]): diff = y_pred[:n_eval, target_label] - t[:n_eval, target_label] print('target_label = {}, y_pred = {}, t = {}, diff = {}'.format( target_label, y_pred[:n_eval, target_label], 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)() # Prevents the loss function from becoming a cupy.core.core.ndarray object # when using the GPU. This hack will be removed as soon as the cause of # the issue is found and properly fixed. loss = numpy.asscalar(cuda.to_cpu(eval_result['main/loss'])) eval_result['main/loss'] = loss print('Evaluation result: ', eval_result) # Save the evaluation results. with open(os.path.join(args.in_dir, 'eval_result.json'), 'w') as f: json.dump(eval_result, f)
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('Fit StandardScaler to the labels.') scaler = StandardScaler() scaler.fit(dataset.get_datasets()[-1]) 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 regressor. device = chainer.get_device(args.device) metrics_fun = {'mae': F.mean_absolute_error, 'rmse': rmse} regressor = Regressor(predictor, lossfun=F.mean_squared_error, metrics_fun=metrics_fun, device=device) print('Training...') run_train(regressor, train, valid=valid, batch_size=args.batchsize, epoch=args.epoch, out=args.out, extensions_list=None, device=device, converter=concat_mols, resume_path=None) # 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 train(): parser = argparser.get_parser() args = parser.parse_args() device = -1 if args.gpu >= 0: device = args.gpu debug = args.debug print('input args:\n', json.dumps(vars(args), indent=4, separators=(',', ':'))) # pretty print args if args.data_name == 'qm9': from data import transform_qm9 transform_fn = transform_qm9.transform_fn atomic_num_list = [6, 7, 8, 9, 0] mlp_channels = [256, 256] gnn_channels = {'gcn': [8, 64], 'hidden': [128, 64]} valid_idx = transform_qm9.get_val_ids() elif args.data_name == 'zinc250k': transform_fn = transform_fn_zinc250k atomic_num_list = zinc250_atomic_num_list mlp_channels = [1024, 512] gnn_channels = {'gcn': [16, 128], 'hidden': [256, 64]} valid_idx = transform_zinc250k.get_val_ids() dataset = NumpyTupleDataset.load( os.path.join(args.data_dir, args.data_file)) dataset = TransformDataset(dataset, transform_fn) if len(valid_idx) > 0: train_idx = [t for t in range(len(dataset)) if t not in valid_idx] n_train = len(train_idx) train_idx.extend(valid_idx) train, test = chainer.datasets.split_dataset(dataset, n_train, train_idx) else: train, test = chainer.datasets.split_dataset_random( dataset, int(len(dataset) * 0.8), seed=args.seed) train_iter = chainer.iterators.SerialIterator(train, args.batch_size) num_masks = { 'node': args.num_node_masks, 'channel': args.num_channel_masks } mask_size = { 'node': args.node_mask_size, 'channel': args.channel_mask_size } num_coupling = { 'node': args.num_node_coupling, 'channel': args.num_channel_coupling } model_params = Hyperparameters( args.num_atoms, args.num_rels, len(atomic_num_list), num_masks=num_masks, mask_size=mask_size, num_coupling=num_coupling, batch_norm=args.apply_batch_norm, additive_transformations=args.additive_transformations, learn_dist=args.learn_dist, mlp_channels=mlp_channels, gnn_channels=gnn_channels) model = GraphNvpModel(model_params) if device >= 0: chainer.cuda.get_device(device).use() model.to_gpu(device) print('==========================================') if device >= 0: print('Using GPUs') print('Num Minibatch-size: {}'.format(args.batch_size)) print('Num epoch: {}'.format(args.max_epochs)) print('==========================================') os.makedirs(args.save_dir, exist_ok=True) model.save_hyperparams(os.path.join(args.save_dir, 'graphnvp-params.json')) opt = chainer.optimizers.Adam() opt.setup(model) updater = MolNvpUpdater(train_iter, opt, device=device, loss_func=None) trainer = training.Trainer(updater, (args.max_epochs, 'epoch'), out=args.save_dir) # trainer.extend(extensions.dump_graph('log_likelihood')) def print_validity(t): adj, x = generate_mols(model, batch_size=100, gpu=device) valid_mols = check_validity(adj, x, atomic_num_list, device)['valid_mols'] mol_dir = os.path.join(args.save_dir, 'generated_{}'.format(t.updater.epoch)) # mol_dir = os.path.join(args.save_dir, 'generated_{}'.format(t.updater.iteration)) os.makedirs(mol_dir, exist_ok=True) for ind, mol in enumerate(valid_mols): save_mol_png(mol, os.path.join(mol_dir, '{}.png'.format(ind))) if debug: # trainer.extend(print_validity, trigger=(1, 'epoch')) trainer.extend(print_validity, trigger=(100, 'iteration')) save_epochs = args.save_epochs if save_epochs == -1: save_epochs = args.max_epochs trainer.extend(extensions.snapshot(), trigger=(save_epochs, 'epoch')) # trainer.extend(extensions.PlotReport(['log_likelihood'], 'epoch', file_name='qm9.png'), # trigger=(100, 'iteration')) trainer.extend( extensions.PrintReport( ['epoch', 'log_likelihood', 'nll_x', 'nll_adj', 'elapsed_time'])) trainer.extend(extensions.LogReport()) trainer.extend(extensions.ProgressBar()) if args.load_params == 1: chainer.serializers.load_npz(args.load_snapshot, trainer) trainer.run() chainer.serializers.save_npz( os.path.join(args.save_dir, 'graph-nvp-final.npz'), model)
def main(): args = parse_arguments() # Set up some useful variables that will be used later on. dataset_name = args.dataset method = args.method num_data = args.num_data if args.label: labels = args.label cache_dir = os.path.join( 'input', '{}_{}_{}'.format(dataset_name, method, labels)) else: labels = None cache_dir = os.path.join('input', '{}_{}_all'.format(dataset_name, method)) # Load the cached dataset. filename = dataset_part_filename('test', num_data) path = os.path.join(cache_dir, filename) if os.path.exists(path): print('Loading cached dataset from {}.'.format(path)) test = NumpyTupleDataset.load(path) else: _, _, test = download_entire_dataset(dataset_name, num_data, labels, method, cache_dir) # Model-related data is stored this directory. model_dir = os.path.join(args.in_dir, os.path.basename(cache_dir)) model_filename = { 'classification': 'classifier.pkl', 'regression': 'regressor.pkl' } task_type = molnet_default_config[dataset_name]['task_type'] model_path = os.path.join(model_dir, model_filename[task_type]) print("model_path=" + model_path) print('Loading model weights from {}...'.format(model_path)) if task_type == 'classification': model = Classifier.load_pickle(model_path, device=args.gpu) elif task_type == 'regression': model = Regressor.load_pickle(model_path, device=args.gpu) else: raise ValueError('Invalid task type ({}) encountered when processing ' 'dataset ({}).'.format(task_type, dataset_name)) # Re-load the best-validation score snapshot # serializers.load_npz(os.path.join( # model_dir, "best_val_" + model_filename[task_type]), model) # Run an evaluator on the test dataset. print('Evaluating...') test_iterator = SerialIterator(test, 16, repeat=False, shuffle=False) eval_result = Evaluator(test_iterator, model, converter=concat_mols, device=args.gpu)() print('Evaluation result: ', eval_result) # Add more stats if task_type == 'regression': # loss = cuda.to_cpu(numpy.array(eval_result['main/loss'])) # eval_result['main/loss'] = loss # convert to native values.. for k, v in eval_result.items(): eval_result[k] = float(v) elif task_type == "classification": # For Classifier, we do not equip the model with ROC-AUC evalation function # use a seperate ROC-AUC Evaluator here rocauc_result = ROCAUCEvaluator(test_iterator, model, converter=concat_mols, device=args.gpu, eval_func=model.predictor, name='test', ignore_labels=-1)() print('ROCAUC Evaluation result: ', rocauc_result) save_json(os.path.join(model_dir, 'rocauc_result.json'), rocauc_result) else: print('[WARNING] unknown task_type {}.'.format(task_type)) # Save the evaluation results. save_json(os.path.join(model_dir, 'eval_result.json'), eval_result)
def main(): args = parse_arguments() # Set up some useful variables that will be used later on. dataset_name = args.dataset method = args.method num_data = args.num_data if args.label: labels = args.label cache_dir = os.path.join( 'input', '{}_{}_{}'.format(dataset_name, method, labels)) else: labels = None cache_dir = os.path.join('input', '{}_{}_all'.format(dataset_name, method)) # Load the cached dataset. filename = dataset_part_filename('test', num_data) path = os.path.join(cache_dir, filename) if os.path.exists(path): print('Loading cached dataset from {}.'.format(path)) test = NumpyTupleDataset.load(path) else: _, _, test = download_entire_dataset(dataset_name, num_data, labels, method, cache_dir) # # Load the standard scaler parameters, if necessary. # if args.scale == 'standardize': # scaler_path = os.path.join(args.in_dir, 'scaler.pkl') # print('Loading scaler parameters from {}.'.format(scaler_path)) # with open(scaler_path, mode='rb') as f: # scaler = pickle.load(f) # else: # print('No standard scaling was selected.') # scaler = None # Model-related data is stored this directory. model_dir = os.path.join(args.in_dir, os.path.basename(cache_dir)) model_filename = { 'classification': 'classifier.pkl', 'regression': 'regressor.pkl' } task_type = molnet_default_config[dataset_name]['task_type'] model_path = os.path.join(model_dir, model_filename[task_type]) print('Loading model weights from {}...'.format(model_path)) if task_type == 'classification': model = Classifier.load_pickle(model_path, device=args.gpu) elif task_type == 'regression': model = Regressor.load_pickle(model_path, device=args.gpu) else: raise ValueError('Invalid task type ({}) encountered when processing ' 'dataset ({}).'.format(task_type, dataset_name)) # # Replace the default predictor with one that scales the output labels. # scaled_predictor = ScaledGraphConvPredictor(model.predictor) # scaled_predictor.scaler = scaler # model.predictor = scaled_predictor # Run an evaluator on the test dataset. print('Evaluating...') test_iterator = SerialIterator(test, 16, repeat=False, shuffle=False) eval_result = Evaluator(test_iterator, model, converter=concat_mols, device=args.gpu)() print('Evaluation result: ', eval_result) # Save the evaluation results. with open(os.path.join(model_dir, 'eval_result.json'), 'w') as f: json.dump(eval_result, f)
def main(): # Parse the arguments. args = parse_arguments() method = args.method if args.data_name == 'suzuki': datafile = 'data/suzuki_type_train_v2.csv' class_num = 119 class_dict = {'M': 28, 'L': 23, 'B': 35, 'S': 10, 'A': 17} dataset_filename = 'data.npz' labels = ['Yield', 'M', 'L', 'B', 'S', 'A', 'id'] elif args.data_name == 'CN': datafile = 'data/CN_coupling_train.csv' class_num = 206 class_dict = {'M': 44, 'L': 47, 'B': 13, 'S': 22, 'A': 74} dataset_filename = 'CN_data.npz' labels = ['Yield', 'M', 'L', 'B', 'S', 'A', 'id'] elif args.data_name == 'Negishi': datafile = 'data/Negishi_train.csv' class_num = 106 class_dict = {'M': 32, 'L': 20, 'T': 8, 'S': 10, 'A': 30} dataset_filename = 'Negishi_data.npz' labels = ['Yield', 'M', 'L', 'T', 'S', 'A', 'id'] elif args.data_name == 'PKR': datafile = 'data/PKR_train.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 = '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. Postprocessing is required for the regression task. def postprocess_label(label_list): return numpy.asarray(label_list, dtype=numpy.float32) # 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...') 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) # Load the entire dataset. 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) # Scale the label values, if necessary. if args.scale == 'standardize': scaler = StandardScaler() labels = dataset.get_datasets()[-2] # yields = dataset.get_datasets()[-1] yields = dataset.get_datasets()[-1][:, 0].reshape(-1, 1).astype('float32') # Filter index here # range_exp = (0.0 <= yields) & (yields <= 1.0) range_exp = numpy.argsort(yields[:, 0]) # ascending start_len = 0 end_len = len(yields) #int(len(yields) / 4) range_exp = range_exp[start_len:end_len] range_dataset = (dataset.get_datasets()[0][range_exp], dataset.get_datasets()[1][range_exp], dataset.get_datasets()[2][range_exp], dataset.get_datasets()[3][range_exp], dataset.get_datasets()[4][range_exp], dataset.get_datasets()[5][range_exp]) yields = yields[range_exp] labels = labels[range_exp] dataset = NumpyTupleDataset(*(range_dataset + ( yields, labels, ))) else: 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(args.method, args.unit_num, args.conv_layers, class_num) # Set up the iterator. 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 classifier = Classifier(predictor, lossfun=F.sigmoid_cross_entropy, metrics_fun=F.binary_accuracy, device=args.gpu) # Set up the optimizer. optimizer = optimizers.Adam() optimizer.setup(classifier) # 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.Evaluator(valid_iter, classifier, device=device, converter=concat_mols)) trainer.extend(E.snapshot(), trigger=(args.epoch, 'epoch')) trainer.extend( E.snapshot_object( classifier, filename='model_epoch-{.updater.epoch}')) # save every 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() # Save the regressor's parameters. model_path = os.path.join(args.out, args.model_filename) print('Saving the trained model to {}...'.format(model_path)) classifier.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)
default='true', help='if neighborhood of a molecule to be visualized') parser.add_argument('--save_fig', type=strtobool, default='true') args = parser.parse_args() chainer.config.train = False snapshot_path = os.path.join(args.model_dir, args.snapshot_path) hyperparams_path = os.path.join(args.model_dir, args.hyperparams_path) print("loading hyperparamaters from {}".format(hyperparams_path)) model_params = Hyperparameters(path=hyperparams_path) model = load_model(snapshot_path, model_params, debug=True) if args.gpu >= 0: model.to_gpu(args.gpu) true_data = NumpyTupleDataset.load( os.path.join(args.data_dir, args.molecule_file)) if args.data_name == 'qm9': atomic_num_list = [6, 7, 8, 9, 0] true_data = TransformDataset(true_data, transform_qm9.transform_fn) valid_idx = transform_qm9.get_val_ids() elif args.data_name == 'zinc250k': atomic_num_list = zinc250_atomic_num_list true_data = TransformDataset(true_data, transform_fn_zinc250k) valid_idx = transform_zinc250k.get_val_ids() train_idx = [t for t in range(len(true_data)) if t not in valid_idx] n_train = len(train_idx) train_idx.extend(valid_idx) train_data, _ = chainer.datasets.split_dataset(true_data, n_train, train_idx)
def main(): args.cuda = torch.cuda.is_available() args.device = 'cuda' if torch.cuda.is_available() else 'cpu' print(args) if args.data_name == 'qm9': from data import transform_qm9 transform_fn = transform_qm9.transform_fn args.atomic_num_list = [6, 7, 8, 9, 0] mlp_channels = [256, 256] gnn_channels = {'gcn': [8, 64], 'hidden': [64, 128]} valid_idx = transform_qm9.get_val_ids() elif args.data_name == 'zinc250k': from data import transform_zinc250k from data.transform_zinc250k import transform_fn_zinc250k transform_fn = transform_fn_zinc250k args.atomic_num_list = [6, 7, 8, 9, 15, 16, 17, 35, 53, 0] mlp_channels = [1024, 512] gnn_channels = {'gcn': [16, 128], 'hidden': [64, 256]} valid_idx = transform_zinc250k.get_val_ids() dataset = NumpyTupleDataset.load(joinpath(args.data_dir, args.data_file)) dataset = TransformDataset(dataset, transform_fn) if len(valid_idx) > 0: train_idx = [t for t in range(len(dataset)) if t not in valid_idx] n_train = len(train_idx) train_idx.extend(valid_idx) train, test = chainer.datasets.split_dataset(dataset, n_train, train_idx) else: train, test = chainer.datasets.split_dataset_random( dataset, int(len(dataset) * 0.8), seed=args.seed) num_masks = { 'node': args.num_node_masks, 'channel': args.num_channel_masks } mask_size = { 'node': args.node_mask_size, 'channel': args.channel_mask_size } num_coupling = { 'node': args.num_node_coupling, 'channel': args.num_channel_coupling } NVPmodel_params = Hyperparameters( args.num_atoms, args.num_rels, len(args.atomic_num_list), num_masks=num_masks, mask_size=mask_size, num_coupling=num_coupling, batch_norm=args.apply_batch_norm, additive_transformations=args.additive_transformations, mlp_channels=mlp_channels, gnn_channels=gnn_channels) model = GraphNvpModel(NVPmodel_params).to(args.device) train_iter = chainer.iterators.SerialIterator(train, args.batch_size, repeat=False) if isinstance(train_iter, iterator_module.Iterator): iterator = {'main': train_iter} # train_dataloader dataloader = iterator['main'] data_config, all_train_smiles = read_molecules(args.data_dir) converter = convert.concat_examples # fitting t_total = time() total_g_loss, total_d_loss = [], [] max_size = model.num_atoms # 9 for QM9 num_atom = max_size node_dim = model.num_features # 5 for QM9 # OR exclude padding dim. 5-1 bond_dim = model.num_bonds # 4 for QM9 best_g_loss, best_d_loss = sys.maxsize, sys.maxsize start_epoch = args.resume_epoch if args.resume: model = GraphNvpModel(hyperparams=NVPmodel_params).to(args.device) model_path = joinpath(args.model_save_dir, 'epoch{}-mle.ckpt'.format(args.resume_epoch)) model.load_state_dict( torch.load(model_path, map_location=lambda storage, loc: storage)) print("Resuming from epoch{}".format(args.resume_epoch)) all_unique_rate = [] all_valid_rate = [] all_novelty_rate = [] print('start fitting.') optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), args.mle_lr, betas=(args.beta1, args.beta2)) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, mode='min', factor=args.lr_decay_factor, patience=args.lr_decay_patience, min_lr=args.lr_decay_min) optimizer.step() def generate_one(model, mute=False, cnt=None): """ inverse flow to generate one molecule Args: temp: temperature of normal distributions, we sample from (0, temp^2 * I) """ generate_start_t = time() num2bond = { 0: Chem.rdchem.BondType.SINGLE, 1: Chem.rdchem.BondType.DOUBLE, 2: Chem.rdchem.BondType.TRIPLE } num2bond_symbol = {0: '=', 1: '==', 2: '==='} num2atom = {0: 6, 1: 7, 2: 8, 3: 9, 4: 15, 5: 16, 6: 17, 7: 35, 8: 53} num2symbol = { 0: 'C', 1: 'N', 2: 'O', 3: 'F', 4: 'P', 5: 'S', 6: 'Cl', 7: 'Br', 8: 'I' } is_continue = True mol = None total_resample = 0 batch_size = 1 # Generating z = sample_z(model, batch_size=1) A, X = model.reverse( z, model.x_size) # For QM9: [16,9,9,5], [16,9,5], [16,8]-[B,z_dim] X = F.softmax(X, dim=2) mols = [ construct_mol(x_elem, adj_elem, args.atomic_num_list) for x_elem, adj_elem in zip(X, A) ] pure_valid = 0 smiles = '' num_atoms = -1 for mol in mols: assert mol is not None, 'mol is None...' final_valid = env.check_chemical_validity(mol) valency_valid = env.check_valency(mol) if final_valid is False or valency_valid is False: print( 'Warning: use valency check during generation but the final molecule is invalid!!!' ) continue num_atoms = mol.GetNumAtoms() num_bonds = mol.GetNumBonds() smiles = Chem.MolToSmiles(mol) if total_resample == 0: pure_valid = 1.0 if not mute: cnt = str(cnt) if cnt is not None else '' print( 'smiles%s: %s | #atoms: %d | #bonds: %d | #resample: %.5f | time: %.5f |' % (cnt, smiles, num_atoms, num_bonds, total_resample, time() - generate_start_t)) return smiles, A, X, pure_valid, num_atoms def train(model): for epoch in range(1 + start_epoch, args.epochs + 1 - start_epoch): batch_g_losses = [] batch_cnt = 0 epoch_example = 0 num_samples = len(dataloader.dataset) num_batches = math.ceil(num_samples / args.batch_size) pbar = tqdm(total=num_batches) for i_batch, batch_data in enumerate(copy.copy(dataloader)): batch_time_s = time() loss = {} in_arrays = converter(batch_data) X, A, label = in_arrays[0], in_arrays[1], in_arrays[2] X, A, label = torch.tensor(X, dtype=torch.float32).to(args.device), \ torch.tensor(A, dtype=torch.float32).to(args.device), \ torch.tensor(label, dtype=torch.float32).to(args.device) # Dequantization X_prime = X + 0.9 * torch.rand(X.shape, device=args.device) A_prime = A + 0.9 * torch.rand(A.shape, device=args.device) z, sum_log_det_jacs = model(A_prime, X_prime) nll = model.log_prob(z, sum_log_det_jacs) g_loss = nll loss['G/loss_g'] = g_loss.item() batch_g_losses.append(g_loss.item()) optimizer.zero_grad() g_loss.backward() optimizer.step() scheduler.step(g_loss) pbar.update() if i_batch % args.show_loss_step == 0: tqdm.write("Epoch %d, batch %d, Loss mle: %.5f" % (epoch, i_batch, g_loss.item())) pbar.close() print("Saving GraphNVP model trained with maximum liklihood") model_path = joinpath(args.model_save_dir, 'epoch{}-mle.ckpt'.format(epoch)) torch.save(model.state_dict(), model_path) print('Saved model checkpoints into {}...'.format( args.model_save_dir)) gen(model, epoch) def gen(model, epoch=-1): model.eval() all_smiles = [] pure_valids = [] appear_in_train = 0. start_t = time() cnt_mol = 0 cnt_gen = 0 out_path = joinpath(args.gen_path, 'mle_mols{}.txt'.format(epoch)) print("Generating %d mols for evaluation" % (args.num_gen)) while cnt_mol < args.num_gen: smiles, A, X, no_resample, num_atoms = generate_one(model, mute=False, cnt=cnt_gen) cnt_gen += 1 if cnt_gen > args.max_resample: break if num_atoms < 0 or num_atoms < args.min_atoms: print('#atoms of generated molecule less than %d, discarded!' % args.min_atoms) continue else: cnt_mol += 1 if cnt_mol % 100 == 0: print('cur cnt mol: %d' % cnt_mol) all_smiles.append(smiles) pure_valids.append(no_resample) print('Accepting: {}'.format(smiles)) if all_train_smiles is not None and smiles in all_train_smiles: appear_in_train += 1.0 mol = Chem.MolFromSmiles(smiles) qed_score = env.qed(mol) plogp_score = env.penalized_logp(mol) if cnt_mol > args.num_gen: print("Generating {} times rather than 100 times!".format(cnt_mol)) args.num_gen = cnt_mol unique_smiles = list(set(all_smiles)) unique_rate = len(unique_smiles) / args.num_gen pure_valid_rate = sum(pure_valids) / args.num_gen novelty = 1. - (appear_in_train / args.num_gen) print( 'Time for generating (%d/%d) molecules(#atoms>=%d) with %d resamplings: %.5f' % (cnt_gen - args.max_resample, args.num_gen, args.min_atoms, args.max_resample, time() - start_t)) print('| unique rate: %.5f | valid rate: %.5f | novelty: %.5f |' % (unique_rate, pure_valid_rate, novelty)) mol_img_dir = joinpath(args.img_dir, 'mol_img{}'.format(epoch)) os.makedirs(mol_img_dir, exist_ok=True) if not os.path.exists(args.gen_path): os.makedirs(args.gen_path) if out_path is not None and args.save: with open(out_path, 'w+') as out_file: cnt = 0 for i, mol in enumerate(all_smiles): # Invalid disconnection if '.' in all_smiles[i]: continue out_file.write(all_smiles[i] + '\n') save_mol_png(Chem.MolFromSmiles(mol), joinpath(mol_img_dir, '{}.png'.format(i))) cnt += 1 print('writing %d smiles into %s done!' % (cnt, out_path)) all_unique_rate.append(unique_rate) all_valid_rate.append(pure_valid_rate) all_novelty_rate.append(novelty) if args.save: print('saving metric of validity, novelty and uniqueness into %s' % (args.gen_path)) np.save(joinpath(args.gen_path, 'valid{}'.format(epoch)), np.array(all_valid_rate)) np.save(joinpath(args.gen_path, 'novelty{}'.format(epoch)), np.array(all_novelty_rate)) np.save(joinpath(args.gen_path, 'unique{}'.format(epoch)), np.array(all_unique_rate)) if args.mode == 'train': train(model) elif args.mode == 'gen': gen(model) else: print("Specify mode as 'train' or 'gen'")
def main(): args = parse_arguments() # Set up some useful variables that will be used later on. dataset_name = args.dataset method = args.method num_data = args.num_data n_unit = args.unit_num conv_layers = args.conv_layers task_type = molnet_default_config[dataset_name]['task_type'] model_filename = { 'classification': 'classifier.pkl', 'regression': 'regressor.pkl' } print('Using dataset: {}...'.format(dataset_name)) # Set up some useful variables that will be used later on. 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']) # Load the train and validation parts of the dataset. filenames = [ dataset_part_filename(p, num_data) for p in ['train', 'valid'] ] paths = [os.path.join(cache_dir, f) for f in filenames] if all([os.path.exists(path) for path in paths]): dataset_parts = [] for path in paths: print('Loading cached dataset from {}.'.format(path)) dataset_parts.append(NumpyTupleDataset.load(path)) else: dataset_parts = download_entire_dataset(dataset_name, num_data, labels, method, cache_dir) train, valid = dataset_parts[0], dataset_parts[1] # # Scale the label values, if necessary. # if args.scale == 'standardize': # if task_type == 'regression': # print('Applying standard scaling to the labels.') # datasets, scaler = standardize_dataset_labels(datasets) # else: # print('Label scaling is not available for classification tasks.') # else: # print('No label scaling was selected.') # scaler = None # Set up the predictor. predictor = set_up_predictor(method, n_unit, conv_layers, class_num) # Set up the iterators. train_iter = iterators.SerialIterator(train, args.batchsize) valid_iter = iterators.SerialIterator(valid, args.batchsize, repeat=False, shuffle=False) # Load metrics for the current dataset. metrics = molnet_default_config[dataset_name]['metrics'] metrics_fun = { k: v for k, v in metrics.items() if isinstance(v, types.FunctionType) } loss_fun = molnet_default_config[dataset_name]['loss'] if task_type == 'regression': model = Regressor(predictor, lossfun=loss_fun, metrics_fun=metrics_fun, device=args.gpu) # TODO: 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 ValueError('Invalid task type ({}) encountered when processing ' 'dataset ({}).'.format(task_type, dataset_name)) # Set up the optimizer. optimizer = optimizers.Adam() optimizer.setup(model) # Save model-related output to this directory. model_dir = os.path.join(args.out, os.path.basename(cache_dir)) if not os.path.exists(model_dir): os.makedirs(model_dir) # Set up the updater. updater = training.StandardUpdater(train_iter, optimizer, device=args.gpu, converter=concat_mols) # Set up the trainer. trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=model_dir) trainer.extend( E.Evaluator(valid_iter, model, device=args.gpu, converter=concat_mols)) trainer.extend(E.snapshot(), trigger=(args.epoch, 'epoch')) trainer.extend(E.LogReport()) # Report various metrics. print_report_targets = ['epoch', 'main/loss', 'validation/main/loss'] for metric_name, metric_fun in metrics.items(): if isinstance(metric_fun, types.FunctionType): print_report_targets.append('main/' + metric_name) print_report_targets.append('validation/main/' + metric_name) elif issubclass(metric_fun, BatchEvaluator): trainer.extend( metric_fun(valid_iter, model, device=args.gpu, eval_func=predictor, converter=concat_mols, name='val', raise_value_error=False)) print_report_targets.append('val/main/' + metric_name) else: raise TypeError('{} is not a supported metrics function.'.format( type(metrics_fun))) print_report_targets.append('elapsed_time') trainer.extend(E.PrintReport(print_report_targets)) trainer.extend(E.ProgressBar()) trainer.run() # Save the model's parameters. model_path = os.path.join(model_dir, model_filename[task_type]) print('Saving the trained model to {}...'.format(model_path)) model.save_pickle(model_path, 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)
def main(): # Parse the arguments. args = parse_arguments() args.out = os.path.join(args.out, args.method) save_args(args, args.out) 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_float(label_list): return numpy.asarray(label_list, dtype=numpy.float32) def postprocess_label_int(label_list): return numpy.asarray(label_list, dtype=numpy.int64) # Apply a preprocessor to the dataset. if args.train: ## training data fn,ext = os.path.splitext(args.train) if ext==".npz": print('Loading training dataset...') train = NumpyTupleDataset.load(args.train) else: print('Preprocessing training dataset...') preprocessor = preprocess_method_dict[args.method]() if args.classification: parser = CSVFileParser(preprocessor, postprocess_label=postprocess_label_int,labels=labels, smiles_col='SMILES') else: parser = CSVFileParser(preprocessor, postprocess_label=postprocess_label_float,labels=labels, smiles_col='SMILES') train = parser.parse(args.train)['dataset'] NumpyTupleDataset.save(os.path.join(args.out,os.path.split(fn)[1]), train) # Scale the label values, if necessary. if args.scale == 'standardize': scaler = StandardScaler() scaler.fit(train.get_datasets()[-1]) else: scaler = None ## test data fn,ext = os.path.splitext(args.val) if ext==".npz": print('Loading test dataset...') test = NumpyTupleDataset.load(args.val) else: print('Preprocessing test dataset...') preprocessor = preprocess_method_dict[args.method]() if args.classification: parser = CSVFileParser(preprocessor, postprocess_label=postprocess_label_int,labels=labels, smiles_col='SMILES') else: parser = CSVFileParser(preprocessor, postprocess_label=postprocess_label_float,labels=labels, smiles_col='SMILES') test = parser.parse(args.val)['dataset'] NumpyTupleDataset.save(os.path.join(args.out,os.path.split(fn)[1]), test) # Set up the model. device = chainer.get_device(args.device) converter = converter_method_dict[args.method] metrics_fun = {'mae': F.mean_absolute_error, 'rmse': rmse} if args.classification: if args.load_model: model = Classifier.load_pickle(args.load_model, device=device) print("model file loaded: ",args.load_model) else: predictor = set_up_predictor(args.method, args.unit_num, args.conv_layers, class_num) model = Classifier(predictor, lossfun=F.sigmoid_cross_entropy, metrics_fun=F.binary_accuracy, device=device) else: if args.load_model: model = Regressor.load_pickle(args.load_model, device=device) print("model file loaded: ",args.load_model) else: predictor = set_up_predictor( args.method+args.method_suffix, args.unit_num, args.conv_layers, class_num, label_scaler=scaler) model = Regressor(predictor, lossfun=F.mean_squared_error, metrics_fun=metrics_fun, device=device) if args.train: if args.balanced_iter: train = BalancedSerialIterator(train, args.batchsize, train.features[:, -1], ignore_labels=-1) train.show_label_stats() print('Training...') log_keys = ['main/mae','main/rmse','validation/main/mae','validation/main/rmse','validation/main/roc_auc'] extensions_list = [extensions.PlotReport(log_keys, 'iteration', trigger=(100, 'iteration'), file_name='loss.png')] if args.eval_roc and args.classification: extensions_list.append(ROCAUCEvaluator( test, model, eval_func=predictor, device=device, converter=converter, name='validation', pos_labels=1, ignore_labels=-1, raise_value_error=False)) save_json(os.path.join(args.out, 'args.json'), vars(args)) run_train(model, train, valid=test, batch_size=args.batchsize, epoch=args.epoch, out=args.out, extensions_list=extensions_list, device=device, converter=converter) #, resume_path=args.resume) # Save the model's parameters. model_path = os.path.join(args.out, args.model_filename) print('Saving the trained model to {}...'.format(model_path)) if hasattr(model.predictor.graph_conv, 'reset_state'): model.predictor.graph_conv.reset_state() model.save_pickle(model_path, protocol=args.protocol) ## prediction it = SerialIterator(test, args.batchsize, repeat=False, shuffle=False) result = [] for batch in it: in_arrays = convert._call_converter(converter, batch, device) with chainer.using_config('train', False), chainer.function.no_backprop_mode(): if isinstance(in_arrays, tuple): res = model(*in_arrays) elif isinstance(in_arrays, dict): res = model(**in_arrays) else: res = model(in_arrays) result.extend(model.y.array.get()) numpy.savetxt(os.path.join(args.out,"result.csv"), numpy.array(result)) eval_result = Evaluator(it, model, converter=converter,device=device)() print('Evaluation result: ', 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('--batchsize', '-b', type=int, default=32) parser.add_argument('--gpu', '-g', type=int, default=-1) parser.add_argument('--in-dir', '-i', type=str, default='result') parser.add_argument('--seed', '-s', type=int, default=777) parser.add_argument('--train-data-ratio', '-t', type=float, default=0.7) 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 = D.get_qm9_label_names() cache_dir = os.path.join('input', '{}_all'.format(method)) # class_num = len(labels) # 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]() dataset = D.get_qm9(preprocessor, labels=labels) if not os.path.exists(cache_dir): os.mkdir(cache_dir) NumpyTupleDataset.save(dataset_cache_path, dataset) if args.scale == 'standardize': # Standard Scaler for labels with open(os.path.join(args.in_dir, 'ss.pkl'), mode='rb') as f: ss = pickle.load(f) else: ss = None train_data_size = int(len(dataset) * train_data_ratio) train, val = split_dataset_random(dataset, train_data_size, seed) regressor = Regressor.load_pickle(os.path.join(args.in_dir, args.model_filename), device=args.gpu) # type: Regressor # We need to feed only input features `x` to `predict`/`predict_proba`. # This converter extracts only inputs (x1, x2, ...) from the features which # consist of input `x` and label `t` (x1, x2, ..., t). def extract_inputs(batch, device=None): return concat_mols(batch, device=device)[:-1] def postprocess_fn(x): if ss is not None: # Model's output is scaled by StandardScaler, # so we need to rescale back. if isinstance(x, Variable): x = x.data scaled_x = ss.inverse_transform(cuda.to_cpu(x)) return scaled_x else: return x print('Predicting...') y_pred = regressor.predict(val, converter=extract_inputs, postprocess_fn=postprocess_fn) print('y_pred.shape = {}, y_pred[:5, 0] = {}'.format( y_pred.shape, y_pred[:5, 0])) t = concat_mols(val, device=-1)[-1] n_eval = 10 # Construct dataframe df_dict = {} for i, l in enumerate(labels): df_dict.update({ 'y_pred_{}'.format(l): y_pred[:, i], 't_{}'.format(l): t[:, i], }) df = pandas.DataFrame(df_dict) # Show random 5 example's prediction/ground truth table print(df.sample(5)) for target_label in range(y_pred.shape[1]): diff = y_pred[:n_eval, target_label] - t[:n_eval, target_label] print('target_label = {}, y_pred = {}, t = {}, diff = {}'.format( target_label, y_pred[:n_eval, target_label], t[:n_eval, target_label], diff)) # --- evaluate --- # To calc loss/accuracy, we can use `Evaluator`, `ROCAUCEvaluator` print('Evaluating...') val_iterator = SerialIterator(val, 16, repeat=False, shuffle=False) eval_result = Evaluator(val_iterator, regressor, converter=concat_mols, device=args.gpu)() print('Evaluation result: ', eval_result)
chainer.config.train = False # -- load model -- # model_hyperparams_path = os.path.join(gen_params.model_root_path, gen_params.model_hyperparams) model_path = os.path.join(gen_params.model_root_path, gen_params.model_file) model_params = Hyperparameter(model_hyperparams_path).subparams("model") model = load_model_from(model_path, model_params) if device >= 0: model.to_gpu(device) # -- load dataset -- # atomic_num_list = get_atomic_num_id(gen_params.atom_id_to_atomic_num) validation_idxs = get_validation_idxs(gen_params.train_validation_split) dataset = NumpyTupleDataset.load( os.path.join(gen_params.data_root_path, gen_params.dataset)) train_idxs = [i for i in range(len(dataset)) if i not in validation_idxs] n_train = len(train_idxs) train_idxs.extend(validation_idxs) train_set = chainer.datasets.SubDataset(dataset, 0, n_train, order=train_idxs) train_smiles = adj_to_smiles(train_set, atomic_num_list) # -- random generation -- # valid_ratio = [] unique_ratio = [] novel_ratio = [] if save_fig: gen_dir = os.path.join(gen_params.output_root_path, "generated")
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 main(): args = parse_arguments() # Set up some useful variables that will be used later on. dataset_name = args.dataset method = args.method num_data = args.num_data if args.label: labels = args.label cache_dir = os.path.join( 'input', '{}_{}_{}'.format(dataset_name, method, labels)) else: labels = None cache_dir = os.path.join('input', '{}_{}_all'.format(dataset_name, method)) # Load the cached dataset. filename = dataset_part_filename('test', num_data) path = os.path.join(cache_dir, filename) if os.path.exists(path): print('Loading cached dataset from {}.'.format(path)) test = NumpyTupleDataset.load(path) else: _, _, test = download_entire_dataset(dataset_name, num_data, labels, method, cache_dir) # # Load the standard scaler parameters, if necessary. # if args.scale == 'standardize': # scaler_path = os.path.join(args.in_dir, 'scaler.pkl') # print('Loading scaler parameters from {}.'.format(scaler_path)) # with open(scaler_path, mode='rb') as f: # scaler = pickle.load(f) # else: # print('No standard scaling was selected.') # scaler = None # Model-related data is stored this directory. model_dir = os.path.join(args.in_dir, os.path.basename(cache_dir)) model_filename = { 'classification': 'classifier.pkl', 'regression': 'regressor.pkl' } task_type = molnet_default_config[dataset_name]['task_type'] model_path = os.path.join(model_dir, model_filename[task_type]) print("model_path=" + model_path) print('Loading model weights from {}...'.format(model_path)) if task_type == 'classification': model = Classifier.load_pickle(model_path, device=args.gpu) elif task_type == 'regression': model = Regressor.load_pickle(model_path, device=args.gpu) else: raise ValueError('Invalid task type ({}) encountered when processing ' 'dataset ({}).'.format(task_type, dataset_name)) # Proposed by Ishiguro # ToDo: consider go/no-go with following modification # Re-load the best-validation score snapshot serializers.load_npz( os.path.join(model_dir, "best_val_" + model_filename[task_type]), model) # # Replace the default predictor with one that scales the output labels. # scaled_predictor = ScaledGraphConvPredictor(model.predictor) # scaled_predictor.scaler = scaler # model.predictor = scaled_predictor # Run an evaluator on the test dataset. print('Evaluating...') test_iterator = SerialIterator(test, 16, repeat=False, shuffle=False) eval_result = Evaluator(test_iterator, model, converter=concat_mols, device=args.gpu)() print('Evaluation result: ', eval_result) # Proposed by Ishiguro: add more stats # ToDo: considre go/no-go with the following modification if task_type == 'regression': # loss = cuda.to_cpu(numpy.array(eval_result['main/loss'])) # eval_result['main/loss'] = loss # convert to native values.. for k, v in eval_result.items(): eval_result[k] = float(v) save_json(os.path.join(args.in_dir, 'eval_result.json'), eval_result) elif task_type == "classification": # For Classifier, we do not equip the model with ROC-AUC evalation function # use a seperate ROC-AUC Evaluator here rocauc_result = ROCAUCEvaluator(test_iterator, model, converter=concat_mols, device=args.gpu, eval_func=model.predictor, name='test', ignore_labels=-1)() print('ROCAUC Evaluation result: ', rocauc_result) save_json(os.path.join(args.in_dir, 'eval_result.json'), rocauc_result) else: pass # Save the evaluation results. save_json(os.path.join(model_dir, 'eval_result.json'), eval_result)
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(): # 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(): args = parse_arguments() # Set up some useful variables that will be used later on. dataset_name = args.dataset method = args.method num_data = args.num_data n_unit = args.unit_num conv_layers = args.conv_layers task_type = molnet_default_config[dataset_name]['task_type'] model_filename = { 'classification': 'classifier.pkl', 'regression': 'regressor.pkl' } print('Using dataset: {}...'.format(dataset_name)) # Set up some useful variables that will be used later on. 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']) # Load the train and validation parts of the dataset. filenames = [ dataset_part_filename(p, num_data) for p in ['train', 'valid'] ] paths = [os.path.join(cache_dir, f) for f in filenames] if all([os.path.exists(path) for path in paths]): dataset_parts = [] for path in paths: print('Loading cached dataset from {}.'.format(path)) dataset_parts.append(NumpyTupleDataset.load(path)) else: dataset_parts = download_entire_dataset(dataset_name, num_data, labels, method, cache_dir) train, valid = dataset_parts[0], dataset_parts[1] # Scale the label values, if necessary. scaler = None if args.scale == 'standardize': if task_type == 'regression': print('Applying standard scaling to the labels.') scaler = fit_scaler(dataset_parts) else: print('Label scaling is not available for classification tasks.') else: print('No label scaling was selected.') # Set up the predictor. predictor = set_up_predictor(method, n_unit, conv_layers, class_num, label_scaler=scaler) # Set up the iterators. train_iter = iterators.SerialIterator(train, args.batchsize) valid_iter = iterators.SerialIterator(valid, args.batchsize, repeat=False, shuffle=False) # Load metrics for the current dataset. metrics = molnet_default_config[dataset_name]['metrics'] metrics_fun = { k: v for k, v in metrics.items() if isinstance(v, types.FunctionType) } loss_fun = molnet_default_config[dataset_name]['loss'] device = chainer.get_device(args.device) if task_type == 'regression': model = Regressor(predictor, lossfun=loss_fun, metrics_fun=metrics_fun, device=device) elif task_type == 'classification': model = Classifier(predictor, lossfun=loss_fun, metrics_fun=metrics_fun, device=device) else: raise ValueError('Invalid task type ({}) encountered when processing ' 'dataset ({}).'.format(task_type, dataset_name)) # Set up the optimizer. optimizer = optimizers.Adam() optimizer.setup(model) # Save model-related output to this directory. model_dir = os.path.join(args.out, os.path.basename(cache_dir)) if not os.path.exists(model_dir): os.makedirs(model_dir) # 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=model_dir) trainer.extend( E.Evaluator(valid_iter, model, device=device, converter=concat_mols)) trainer.extend(E.snapshot(), trigger=(args.epoch, 'epoch')) trainer.extend(E.LogReport()) # TODO: consider go/no-go of the following block # # (i) more reporting for val/evalutaion # # (ii) best validation score snapshot # if task_type == 'regression': # metric_name_list = list(metrics.keys()) # if 'RMSE' in metric_name_list: # trainer.extend(E.snapshot_object(model, "best_val_" + model_filename[task_type]), # trigger=training.triggers.MinValueTrigger('validation/main/RMSE')) # elif 'MAE' in metric_name_list: # trainer.extend(E.snapshot_object(model, "best_val_" + model_filename[task_type]), # trigger=training.triggers.MinValueTrigger('validation/main/MAE')) # else: # print("[WARNING] No validation metric defined?") # # elif task_type == 'classification': # train_eval_iter = iterators.SerialIterator( # train, args.batchsize, repeat=False, shuffle=False) # trainer.extend(ROCAUCEvaluator( # train_eval_iter, predictor, eval_func=predictor, # device=args.gpu, converter=concat_mols, name='train', # 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( # valid_iter, predictor, eval_func=predictor, # device=args.gpu, converter=concat_mols, name='val', # pos_labels=1, ignore_labels=-1, raise_value_error=False)) # # trainer.extend(E.snapshot_object( # model, "best_val_" + model_filename[task_type]), # trigger=training.triggers.MaxValueTrigger('val/main/roc_auc')) # else: # raise NotImplementedError( # 'Not implemented task_type = {}'.format(task_type)) trainer.extend(AutoPrintReport()) trainer.extend(E.ProgressBar()) trainer.run() # Save the model's parameters. model_path = os.path.join(model_dir, model_filename[task_type]) print('Saving the trained model to {}...'.format(model_path)) model.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) 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() # 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 train(hyperparams: Hyperparameter): # -- hyperparams -- # dataset_params = hyperparams.subparams("dataset") config_params = hyperparams.subparams("configuration") train_params = hyperparams.subparams("train") model_params = hyperparams.subparams("model") output_params = hyperparams.subparams("output") os.makedirs(output_params.root_dir, exist_ok=True) if hasattr(output_params, "logname"): log.basicConfig(filename=os.path.join(output_params.root_dir, output_params.logname), filemode="w", level=get_log_level(output_params.log_level)) else: log.basicConfig(level=get_log_level(output_params.log_level)) hyperparams.save(os.path.join(output_params.root_dir, "hyperparams.json")) atomic_num_list = get_atomic_num_id( os.path.join(config_params.root_dir, config_params.atom_id_to_atomic_num)) data_parallel = False if isinstance(train_params.device, int): main_device = train_params.device device = main_device elif isinstance(train_params.device, dict): main_device = train_params.device["main"] device = train_params.device data_parallel = True else: raise ValueError("Invalid device.") log.info("Main Device: {}".format(main_device)) log.info("dataset hyperparameters:\n{}\n".format(dataset_params)) log.info("configuration hyperparameters:\n{}\n".format(config_params)) log.info("train hyperparameters:\n{}\n".format(train_params)) log.info("model hyperparameters:\n{}\n".format(model_params)) log.info("output hyperparameters:\n{}\n".format(output_params)) # -- build dataset -- # if config_params.has("train_validation_split"): validation_idxs = get_validation_idxs( os.path.join(config_params.root_dir, config_params.train_validation_split)) else: validation_idxs = None dataset = NumpyTupleDataset.load( os.path.join(dataset_params.root_dir, dataset_params.name)) if validation_idxs: train_idxs = [ i for i in range(len(dataset)) if i not in validation_idxs ] trainset_size = len(train_idxs) train_idxs.extend(validation_idxs) trainset, valset = chainer.datasets.split_dataset( dataset, trainset_size, train_idxs) else: trainset, valset = chainer.datasets.split_dataset_random( dataset, int(len(dataset) * 0.8), seed=777) train_iter = chainer.iterators.SerialIterator(trainset, train_params.batch_size, shuffle=True) val_iter = chainer.iterators.SerialIterator(valset, train_params.batch_size, repeat=False, shuffle=False) # -- model -- # model = AttentionNvpModel(model_params) if isinstance(device, dict): log.info("Using multi-GPU {}".format(device)) model.to_gpu(main_device) elif device >= 0: log.info("Using GPU {}".format(device)) chainer.cuda.get_device(main_device).use() model.to_gpu(device) else: log.info("Using CPU") # -- training details -- # num_epoch = train_params.num_epoch opt_gen = get_optimizer(train_params.optimizer) if train_params.has("optimizer_params"): optimizer = opt_gen(**train_params.optimizer_params) else: optimizer = opt_gen() optimizer.setup(model) if data_parallel: updater = DataParallelNVPUpdater( train_iter, optimizer, devices=device, two_step=train_params.two_step, h_nll_weight=train_params.h_nll_weight) else: updater = NVPUpdater(train_iter, optimizer, device=device, two_step=train_params.two_step, h_nll_weight=train_params.h_nll_weight) trainer = training.Trainer(updater, (num_epoch, "epoch"), out=output_params.root_dir) if train_params.has("save_epoch"): save_epoch = train_params.save_epoch else: save_epoch = num_epoch # -- evaluation function -- # def print_validity(trainer): with chainer.using_device( chainer.backends.cuda.get_device_from_id( main_device)), chainer.using_config("train", False): save_mol = (get_log_level(output_params.log_level) <= log.DEBUG) x, adj = generate_mols(model, batch_size=100, device=main_device) # x: atom id valid_mols = check_validity(x, adj, atomic_num_list=atomic_num_list, device=main_device) if save_mol: mol_dir = os.path.join( output_params.root_dir, output_params.saved_mol_dir, "generated_{}".format(trainer.updater.epoch)) os.makedirs(mol_dir, exist_ok=True) for i, mol in enumerate(valid_mols["valid_mols"]): save_mol_png(mol, os.path.join(mol_dir, "{}.png".format(i))) # -- trainer extension -- # trainer.extend(extensions.snapshot(), trigger=(save_epoch, "epoch")) trainer.extend(extensions.LogReport(filename=output_params.trainlogname)) trainer.extend(print_validity, trigger=(1, "epoch")) trainer.extend( extensions.PrintReport([ "epoch", "neg_log_likelihood", "nll_x", "nll_adj", "z_var", "ln_det_x", "ln_det_adj", "elapsed_time" ])) trainer.extend(extensions.ProgressBar()) # -- start train -- # if hasattr(train_params, "load_snapshot"): log.info("Load snapshot from {}".format(train_params.load_snapshot)) chainer.serializers.load_npz(train_params.load_snapshot, trainer) trainer.run() chainer.serializers.save_npz( os.path.join(output_params.root_dir, output_params.final_model_name), model)
def main(): args = parse_arguments() # Set up some useful variables that will be used later on. dataset_name = args.dataset method = args.method num_data = args.num_data n_unit = args.unit_num conv_layers = args.conv_layers task_type = molnet_default_config[dataset_name]['task_type'] model_filename = {'classification': 'classifier.pkl', 'regression': 'regressor.pkl'} print('Using dataset: {}...'.format(dataset_name)) # Set up some useful variables that will be used later on. 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']) # Load the train and validation parts of the dataset. filenames = [dataset_part_filename(p, num_data) for p in ['train', 'valid']] paths = [os.path.join(cache_dir, f) for f in filenames] if all([os.path.exists(path) for path in paths]): dataset_parts = [] for path in paths: print('Loading cached dataset from {}.'.format(path)) dataset_parts.append(NumpyTupleDataset.load(path)) else: dataset_parts = download_entire_dataset(dataset_name, num_data, labels, method, cache_dir) train, valid = dataset_parts[0], dataset_parts[1] # # Scale the label values, if necessary. # if args.scale == 'standardize': # if task_type == 'regression': # print('Applying standard scaling to the labels.') # datasets, scaler = standardize_dataset_labels(datasets) # else: # print('Label scaling is not available for classification tasks.') # else: # print('No label scaling was selected.') # scaler = None # Set up the predictor. predictor = set_up_predictor(method, n_unit, conv_layers, class_num) # Set up the iterators. train_iter = iterators.SerialIterator(train, args.batchsize) valid_iter = iterators.SerialIterator(valid, args.batchsize, repeat=False, shuffle=False) # Load metrics for the current dataset. metrics = molnet_default_config[dataset_name]['metrics'] metrics_fun = {k: v for k, v in metrics.items() if isinstance(v, types.FunctionType)} loss_fun = molnet_default_config[dataset_name]['loss'] if task_type == 'regression': model = Regressor(predictor, lossfun=loss_fun, metrics_fun=metrics_fun, device=args.gpu) # TODO: 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 ValueError('Invalid task type ({}) encountered when processing ' 'dataset ({}).'.format(task_type, dataset_name)) # Set up the optimizer. optimizer = optimizers.Adam() optimizer.setup(model) # Save model-related output to this directory. model_dir = os.path.join(args.out, os.path.basename(cache_dir)) if not os.path.exists(model_dir): os.makedirs(model_dir) # Set up the updater. updater = training.StandardUpdater(train_iter, optimizer, device=args.gpu, converter=concat_mols) # Set up the trainer. trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=model_dir) trainer.extend(E.Evaluator(valid_iter, model, device=args.gpu, converter=concat_mols)) trainer.extend(E.snapshot(), trigger=(args.epoch, 'epoch')) trainer.extend(E.LogReport()) # Report various metrics. print_report_targets = ['epoch', 'main/loss', 'validation/main/loss'] for metric_name, metric_fun in metrics.items(): if isinstance(metric_fun, types.FunctionType): print_report_targets.append('main/' + metric_name) print_report_targets.append('validation/main/' + metric_name) elif issubclass(metric_fun, BatchEvaluator): trainer.extend(metric_fun(valid_iter, model, device=args.gpu, eval_func=predictor, converter=concat_mols, name='val', raise_value_error=False)) print_report_targets.append('val/main/' + metric_name) else: raise TypeError('{} is not a supported metrics function.' .format(type(metrics_fun))) print_report_targets.append('elapsed_time') # Augmented by Ishiguro # ToDo: consider go/no-go of the following block # (i) more reporting for val/evalutaion # (ii) best validation score snapshot if task_type == 'regression': if 'RMSE' in metric_name: trainer.extend(E.snapshot_object(model, "best_val_" + model_filename[task_type]), trigger=training.triggers.MinValueTrigger('validation/main/RMSE')) elif 'MAE' in metric_name: trainer.extend(E.snapshot_object(model, "best_val_" + model_filename[task_type]), trigger=training.triggers.MinValueTrigger('validation/main/MAE')) else: print("No validation metric defined?") assert(False) elif task_type == 'classification': train_eval_iter = iterators.SerialIterator(train, args.batchsize,repeat=False, shuffle=False) trainer.extend(ROCAUCEvaluator( train_eval_iter, predictor, eval_func=predictor, device=args.gpu, converter=concat_mols, name='train', 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( valid_iter, predictor, eval_func=predictor, device=args.gpu, converter=concat_mols, name='val', pos_labels=1, ignore_labels=-1)) print_report_targets.append('train/main/roc_auc') print_report_targets.append('validation/main/loss') print_report_targets.append('val/main/roc_auc') trainer.extend(E.snapshot_object(model, "best_val_" + model_filename[task_type]), trigger=training.triggers.MaxValueTrigger('val/main/roc_auc')) else: raise NotImplementedError( 'Not implemented task_type = {}'.format(task_type)) trainer.extend(E.PrintReport(print_report_targets)) trainer.extend(E.ProgressBar()) trainer.run() # Save the model's parameters. model_path = os.path.join(model_dir, model_filename[task_type]) print('Saving the trained model to {}...'.format(model_path)) model.save_pickle(model_path, protocol=args.protocol)
def train(config): # -- read hyperparameters -- log.info("Hyper-parameters:") device = get_and_log(config, "device", -1) out_dir = get_and_log(config, "out_dir", "./output") config_dir = get_and_log(config, "config_dir", "./config") dataset_dir = get_and_log(config, "dataset_dir", "./dataset") validation_idxs_filepath = get_and_log(config, "train_validation_split") dataset_name = get_and_log(config, "dataset", required=True) atomic_nums = get_and_log(config, "atom_id_to_atomic_num", required=True) batch_size = get_and_log(config, "batch_size", required=True) num_epoch = get_and_log(config, "num_epoch", required=True) word_size = get_and_log(config, "embed_word_size", required=True) molecule_size = get_and_log(config, "molecule_size", required=True) num_atom_type = get_and_log(config, "num_atom_type", required=True) save_epoch = get_and_log(config, "save_epoch", -1) kekulized = get_and_log(config, "kekulize", False) layers = get_and_log(config, "layers", required=True) scale_adj = get_and_log(config, "scale_adj", True) log_name = get_and_log(config, "log_name", "log") optimizer_type = get_and_log(config, "optimizer", "adam") optimizer_params = get_and_log(config, "optimizer_params") snapshot = get_and_log(config, "snapshot") num_edge_type = 4 if kekulized else 5 os.makedirs(out_dir, exist_ok=True) if validation_idxs_filepath is not None: validation_idxs = get_validation_idxs(os.path.join(config_dir, validation_idxs_filepath)) else: validation_idxs = None # -- build dataset -- dataset = NumpyTupleDataset.load(os.path.join(dataset_dir, dataset_name)) if validation_idxs: train_idxs = [i for i in range(len(dataset)) if i not in validation_idxs] trainset_size = len(train_idxs) train_idxs.extend(validation_idxs) trainset, testset = chainer.datasets.split_dataset(dataset, trainset_size, train_idxs) else: trainset, testset = chainer.datasets.split_dataset_random(dataset, int(len(dataset) * 0.8), seed=777) train_iter = chainer.iterators.SerialIterator(trainset, batch_size, shuffle=True) test_iter = chainer.iterators.SerialIterator(testset, batch_size, repeat=False, shuffle=False) # -- model -- model = AtomEmbedModel(word_size, num_atom_type, num_edge_type, layers, scale_adj) model.save_hyperparameters(os.path.join(out_dir, "atom_embed_model_hyper.json")) # -- training details -- if device >= 0: log.info("Using GPU") chainer.cuda.get_device(device).use() model.to_gpu(device) opt_func = get_optimizer(optimizer_type) if optimizer_params is not None: optimizer = opt_func(optimizer_params) else: optimizer = opt_func() optimizer.setup(model) updater = AtomEmbedUpdater(train_iter, optimizer, device=device) trainer = training.Trainer(updater, (num_epoch, "epoch"), out=out_dir) save_epoch = save_epoch if save_epoch >= 0 else num_epoch # -- trainer extension -- trainer.extend(extensions.snapshot, trigger=(save_epoch, "epoch")) trainer.extend(extensions.LogReport(filename=log_name)) trainer.extend(AtomEmbedEvaluator(test_iter, model, reporter=trainer.reporter, device=device)) trainer.extend(extensions.PrintReport(["epoch", "ce_loss", "accuracy", "validation/ce_loss", "validation/accuracy", "elapsed_time"])) trainer.extend(extensions.PlotReport(["ce_loss", "validation/ce_loss"], x_key="epoch", filename="cross_entrypy_loss.png")) trainer.extend(extensions.PlotReport(["accuracy", "validation/accuracy"], x_key="epoch", filename="accuracy.png")) if snapshot is not None: chainer.serializers.load_npz(snapshot, trainer) trainer.run() chainer.serializers.save_npz(os.path.join(out_dir, "final_embed_model.npz"), model)