def load_checkpoint(path: str,
device: torch.device = None,
logger: logging.Logger = None) -> MoleculeModel:
"""
Loads a model checkpoint.
:param path: Path where checkpoint is saved.
:param device: Device where the model will be moved.
:param logger: A logger for recording output.
:return: The loaded :class:`~chemprop.models.model.MoleculeModel`.
"""
if logger is not None:
debug, info = logger.debug, logger.info
else:
debug = info = print
# Load model and args
state = torch.load(path, map_location=lambda storage, loc: storage)
args = TrainArgs()
args.from_dict(vars(state['args']), skip_unsettable=True)
loaded_state_dict = state['state_dict']
if device is not None:
args.device = device
# Build model
model = MoleculeModel(args)
model_state_dict = model.state_dict()
# Skip missing parameters and parameters of mismatched size
pretrained_state_dict = {}
for loaded_param_name in loaded_state_dict.keys():
# Backward compatibility for parameter names
if re.match(r'(encoder\.encoder\.)([Wc])', loaded_param_name):
param_name = loaded_param_name.replace('encoder.encoder', 'encoder.encoder.0')
else:
param_name = loaded_param_name
# Load pretrained parameter, skipping unmatched parameters
if param_name not in model_state_dict:
info(f'Warning: Pretrained parameter "{loaded_param_name}" cannot be found in model parameters.')
elif model_state_dict[param_name].shape != loaded_state_dict[loaded_param_name].shape:
info(f'Warning: Pretrained parameter "{loaded_param_name}" '
f'of shape {loaded_state_dict[loaded_param_name].shape} does not match corresponding '
f'model parameter of shape {model_state_dict[param_name].shape}.')
else:
debug(f'Loading pretrained parameter "{loaded_param_name}".')
pretrained_state_dict[param_name] = loaded_state_dict[loaded_param_name]
# Load pretrained weights
model_state_dict.update(pretrained_state_dict)
model.load_state_dict(model_state_dict)
if args.cuda:
debug('Moving model to cuda')
model = model.to(args.device)
return model
def load_args(path: str) -> TrainArgs:
"""
Loads the arguments a model was trained with.
:param path: Path where model checkpoint is saved.
:return: The :class:`~chemprop.args.TrainArgs` object that the model was trained with.
"""
args = TrainArgs()
args.from_dict(vars(torch.load(path, map_location=lambda storage, loc: storage)['args']), skip_unsettable=True)
return args
def save_checkpoint(path: str,
model: MoleculeModel,
scaler: StandardScaler = None,
features_scaler: StandardScaler = None,
atom_descriptor_scaler: StandardScaler = None,
bond_feature_scaler: StandardScaler = None,
args: TrainArgs = None) -> None:
"""
Saves a model checkpoint.
:param model: A :class:`~chemprop.models.model.MoleculeModel`.
:param scaler: A :class:`~chemprop.data.scaler.StandardScaler` fitted on the data.
:param features_scaler: A :class:`~chemprop.data.scaler.StandardScaler` fitted on the features.
:param atom_descriptor_scaler: A :class:`~chemprop.data.scaler.StandardScaler` fitted on the atom descriptors.
:param bond_feature_scaler: A :class:`~chemprop.data.scaler.StandardScaler` fitted on the bond_fetaures.
:param args: The :class:`~chemprop.args.TrainArgs` object containing the arguments the model was trained with.
:param path: Path where checkpoint will be saved.
"""
# Convert args to namespace for backwards compatibility
if args is not None:
args = Namespace(**args.as_dict())
state = {
'args': args,
'state_dict': model.state_dict(),
'data_scaler': {
'means': scaler.means,
'stds': scaler.stds
} if scaler is not None else None,
'features_scaler': {
'means': features_scaler.means,
'stds': features_scaler.stds
} if features_scaler is not None else None,
'atom_descriptor_scaler': {
'means': atom_descriptor_scaler.means,
'stds': atom_descriptor_scaler.stds
} if atom_descriptor_scaler is not None else None,
'bond_feature_scaler': {
'means': bond_feature_scaler.means,
'stds': bond_feature_scaler.stds
} if bond_feature_scaler is not None else None
}
torch.save(state, path)
def cross_validate(args: TrainArgs,
train_func: Callable[[TrainArgs, MoleculeDataset, Logger], Dict[str, List[float]]]
) -> Tuple[float, float]:
"""
Runs k-fold cross-validation.
For each of k splits (folds) of the data, trains and tests a model on that split
and aggregates the performance across folds.
:param args: A :class:`~chemprop.args.TrainArgs` object containing arguments for
loading data and training the Chemprop model.
:param train_func: Function which runs training.
:return: A tuple containing the mean and standard deviation performance across folds.
"""
logger = create_logger(name=TRAIN_LOGGER_NAME, save_dir=args.save_dir, quiet=args.quiet)
if logger is not None:
debug, info = logger.debug, logger.info
else:
debug = info = print
# Initialize relevant variables
init_seed = args.seed
save_dir = args.save_dir
args.task_names = get_task_names(path=args.data_path, smiles_columns=args.smiles_columns,
target_columns=args.target_columns, ignore_columns=args.ignore_columns)
# Print command line
debug('Command line')
debug(f'python {" ".join(sys.argv)}')
# Print args
debug('Args')
debug(args)
# Save args
makedirs(args.save_dir)
args.save(os.path.join(args.save_dir, 'args.json'))
#set explicit H option and reaction option
set_explicit_h(args.explicit_h)
set_reaction(args.reaction, args.reaction_mode)
# Get data
debug('Loading data')
data = get_data(
path=args.data_path,
args=args,
smiles_columns=args.smiles_columns,
features_columns=args.features_columns,
logger=logger,
skip_none_targets=True
)
validate_dataset_type(data, dataset_type=args.dataset_type)
args.features_size = data.features_size()
if args.atom_descriptors == 'descriptor':
args.atom_descriptors_size = data.atom_descriptors_size()
args.ffn_hidden_size += args.atom_descriptors_size
elif args.atom_descriptors == 'feature':
args.atom_features_size = data.atom_features_size()
set_extra_atom_fdim(args.atom_features_size)
if args.bond_features_path is not None:
args.bond_features_size = data.bond_features_size()
set_extra_bond_fdim(args.bond_features_size)
debug(f'Number of tasks = {args.num_tasks}')
# Run training on different random seeds for each fold
all_scores = defaultdict(list)
for fold_num in range(args.num_folds):
info(f'Fold {fold_num}')
args.seed = init_seed + fold_num
args.save_dir = os.path.join(save_dir, f'fold_{fold_num}')
makedirs(args.save_dir)
data.reset_features_and_targets()
# If resuming experiment, load results from trained models
test_scores_path = os.path.join(args.save_dir, 'test_scores.json')
if args.resume_experiment and os.path.exists(test_scores_path):
print('Loading scores')
with open(test_scores_path) as f:
model_scores = json.load(f)
# Otherwise, train the models
else:
model_scores = train_func(args, data, logger)
for metric, scores in model_scores.items():
all_scores[metric].append(scores)
all_scores = dict(all_scores)
# Convert scores to numpy arrays
for metric, scores in all_scores.items():
all_scores[metric] = np.array(scores)
# Report results
info(f'{args.num_folds}-fold cross validation')
# Report scores for each fold
for fold_num in range(args.num_folds):
for metric, scores in all_scores.items():
info(f'\tSeed {init_seed + fold_num} ==> test {metric} = {np.nanmean(scores[fold_num]):.6f}')
if args.show_individual_scores:
for task_name, score in zip(args.task_names, scores[fold_num]):
info(f'\t\tSeed {init_seed + fold_num} ==> test {task_name} {metric} = {score:.6f}')
# Report scores across folds
for metric, scores in all_scores.items():
avg_scores = np.nanmean(scores, axis=1) # average score for each model across tasks
mean_score, std_score = np.nanmean(avg_scores), np.nanstd(avg_scores)
info(f'Overall test {metric} = {mean_score:.6f} +/- {std_score:.6f}')
if args.show_individual_scores:
for task_num, task_name in enumerate(args.task_names):
info(f'\tOverall test {task_name} {metric} = '
f'{np.nanmean(scores[:, task_num]):.6f} +/- {np.nanstd(scores[:, task_num]):.6f}')
# Save scores
with open(os.path.join(save_dir, TEST_SCORES_FILE_NAME), 'w') as f:
writer = csv.writer(f)
header = ['Task']
for metric in args.metrics:
header += [f'Mean {metric}', f'Standard deviation {metric}'] + \
[f'Fold {i} {metric}' for i in range(args.num_folds)]
writer.writerow(header)
for task_num, task_name in enumerate(args.task_names):
row = [task_name]
for metric, scores in all_scores.items():
task_scores = scores[:, task_num]
mean, std = np.nanmean(task_scores), np.nanstd(task_scores)
row += [mean, std] + task_scores.tolist()
writer.writerow(row)
# Determine mean and std score of main metric
avg_scores = np.nanmean(all_scores[args.metric], axis=1)
mean_score, std_score = np.nanmean(avg_scores), np.nanstd(avg_scores)
# Optionally merge and save test preds
if args.save_preds:
all_preds = pd.concat([pd.read_csv(os.path.join(save_dir, f'fold_{fold_num}', 'test_preds.csv'))
for fold_num in range(args.num_folds)])
all_preds.to_csv(os.path.join(save_dir, 'test_preds.csv'), index=False)
return mean_score, std_score
def chemprop_train() -> None:
"""Parses Chemprop training arguments and trains (cross-validates) a Chemprop model.
This is the entry point for the command line command :code:`chemprop_train`.
"""
cross_validate(args=TrainArgs().parse_args(), train_func=run_training)
def update_prediction_args(predict_args: PredictArgs,
train_args: TrainArgs,
missing_to_defaults: bool = True,
validate_feature_sources: bool = True) -> None:
"""
Updates prediction arguments with training arguments loaded from a checkpoint file.
If an argument is present in both, the prediction argument will be used.
Also raises errors for situations where the prediction arguments and training arguments
are different but must match for proper function.
:param predict_args: The :class:`~chemprop.args.PredictArgs` object containing the arguments to use for making predictions.
:param train_args: The :class:`~chemprop.args.TrainArgs` object containing the arguments used to train the model previously.
:param missing_to_defaults: Whether to replace missing training arguments with the current defaults for :class: `~chemprop.args.TrainArgs`.
This is used for backwards compatibility.
:param validate_feature_sources: Indicates whether the feature sources (from path or generator) are checked for consistency between
the training and prediction arguments. This is not necessary for fingerprint generation, where molecule features are not used.
"""
for key, value in vars(train_args).items():
if not hasattr(predict_args, key):
setattr(predict_args, key, value)
if missing_to_defaults:
# If a default argument would cause different behavior than occurred in legacy checkpoints before the argument existed,
# then that argument must be included in the `override_defaults` dictionary to force the legacy behavior.
override_defaults = {
'bond_features_scaling':False,
'no_bond_features_scaling':True,
'atom_descriptors_scaling':False,
'no_atom_descriptors_scaling':True,
}
default_train_args=TrainArgs().parse_args(['--data_path', None, '--dataset_type', str(train_args.dataset_type)])
for key, value in vars(default_train_args).items():
if not hasattr(predict_args,key):
setattr(predict_args,key,override_defaults.get(key,value))
# Same number of molecules must be used in training as in making predictions
if train_args.number_of_molecules != predict_args.number_of_molecules:
raise ValueError('A different number of molecules was used in training '
f'model than is specified for prediction, {train_args.number_of_molecules} '
'smiles fields must be provided')
# If atom-descriptors were used during training, they must be used when predicting and vice-versa
if train_args.atom_descriptors != predict_args.atom_descriptors:
raise ValueError('The use of atom descriptors is inconsistent between training and prediction. If atom descriptors '
' were used during training, they must be specified again during prediction using the same type of '
' descriptors as before. If they were not used during training, they cannot be specified during prediction.')
# If bond features were used during training, they must be used when predicting and vice-versa
if (train_args.bond_features_path is None) != (predict_args.bond_features_path is None):
raise ValueError('The use of bond descriptors is different between training and prediction. If you used bond '
'descriptors for training, please specify a path to new bond descriptors for prediction.')
# if atom or bond features were scaled, the same must be done during prediction
if train_args.features_scaling != predict_args.features_scaling:
raise ValueError('If scaling of the additional features was done during training, the '
'same must be done during prediction.')
# If atom descriptors were used during training, they must be used when predicting and vice-versa
if train_args.atom_descriptors != predict_args.atom_descriptors:
raise ValueError('The use of atom descriptors is inconsistent between training and prediction. '
'If atom descriptors were used during training, they must be specified again '
'during prediction using the same type of descriptors as before. '
'If they were not used during training, they cannot be specified during prediction.')
# If bond features were used during training, they must be used when predicting and vice-versa
if (train_args.bond_features_path is None) != (predict_args.bond_features_path is None):
raise ValueError('The use of bond descriptors is different between training and prediction. If you used bond'
'descriptors for training, please specify a path to new bond descriptors for prediction.')
if validate_feature_sources:
# If features were used during training, they must be used when predicting
if ((train_args.features_path is None) != (predict_args.features_path is None)):
# or ((train_args.features_generator is None) != (predict_args.features_generator is None))):
raise ValueError('Features were used during training so they must be specified again during prediction '
'using the same type of features as before (with either --features_generator or '
'--features_path and using --no_features_scaling if applicable).')
def run_training(args: TrainArgs,
data: MoleculeDataset,
logger: Logger = None) -> Dict[str, List[float]]:
"""
Loads data, trains a Chemprop model, and returns test scores for the model checkpoint with the highest validation score.
:param args: A :class:`~chemprop.args.TrainArgs` object containing arguments for
loading data and training the Chemprop model.
:param data: A :class:`~chemprop.data.MoleculeDataset` containing the data.
:param logger: A logger to record output.
:return: A dictionary mapping each metric in :code:`args.metrics` to a list of values for each task.
"""
if logger is not None:
debug, info = logger.debug, logger.info
else:
debug = info = print
# Set pytorch seed for random initial weights
torch.manual_seed(args.pytorch_seed)
# Split data
debug(f'Splitting data with seed {args.seed}')
if args.separate_test_path:
test_data = get_data(path=args.separate_test_path,
args=args,
features_path=args.separate_test_features_path,
atom_descriptors_path=args.separate_test_atom_descriptors_path,
bond_features_path=args.separate_test_bond_features_path,
smiles_columns=args.smiles_columns,
logger=logger)
if args.separate_val_path:
val_data = get_data(path=args.separate_val_path,
args=args,
features_path=args.separate_val_features_path,
atom_descriptors_path=args.separate_val_atom_descriptors_path,
bond_features_path=args.separate_val_bond_features_path,
smiles_columns = args.smiles_columns,
logger=logger)
if args.separate_val_path and args.separate_test_path:
train_data = data
elif args.separate_val_path:
train_data, _, test_data = split_data(data=data,
split_type=args.split_type,
sizes=(0.8, 0.0, 0.2),
seed=args.seed,
num_folds=args.num_folds,
args=args,
logger=logger)
elif args.separate_test_path:
train_data, val_data, _ = split_data(data=data,
split_type=args.split_type,
sizes=(0.8, 0.2, 0.0),
seed=args.seed,
num_folds=args.num_folds,
args=args,
logger=logger)
else:
train_data, val_data, test_data = split_data(data=data,
split_type=args.split_type,
sizes=args.split_sizes,
seed=args.seed,
num_folds=args.num_folds,
args=args,
logger=logger)
if args.dataset_type == 'classification':
class_sizes = get_class_sizes(data)
debug('Class sizes')
for i, task_class_sizes in enumerate(class_sizes):
debug(f'{args.task_names[i]} '
f'{", ".join(f"{cls}: {size * 100:.2f}%" for cls, size in enumerate(task_class_sizes))}')
if args.save_smiles_splits:
save_smiles_splits(
data_path=args.data_path,
save_dir=args.save_dir,
task_names=args.task_names,
features_path=args.features_path,
train_data=train_data,
val_data=val_data,
test_data=test_data,
smiles_columns=args.smiles_columns,
logger=logger,
)
if args.features_scaling:
features_scaler = train_data.normalize_features(replace_nan_token=0)
val_data.normalize_features(features_scaler)
test_data.normalize_features(features_scaler)
else:
features_scaler = None
if args.atom_descriptor_scaling and args.atom_descriptors is not None:
atom_descriptor_scaler = train_data.normalize_features(replace_nan_token=0, scale_atom_descriptors=True)
val_data.normalize_features(atom_descriptor_scaler, scale_atom_descriptors=True)
test_data.normalize_features(atom_descriptor_scaler, scale_atom_descriptors=True)
else:
atom_descriptor_scaler = None
if args.bond_feature_scaling and args.bond_features_size > 0:
bond_feature_scaler = train_data.normalize_features(replace_nan_token=0, scale_bond_features=True)
val_data.normalize_features(bond_feature_scaler, scale_bond_features=True)
test_data.normalize_features(bond_feature_scaler, scale_bond_features=True)
else:
bond_feature_scaler = None
args.train_data_size = len(train_data)
debug(f'Total size = {len(data):,} | '
f'train size = {len(train_data):,} | val size = {len(val_data):,} | test size = {len(test_data):,}')
# Initialize scaler and scale training targets by subtracting mean and dividing standard deviation (regression only)
if args.dataset_type == 'regression':
debug('Fitting scaler')
scaler = train_data.normalize_targets()
else:
scaler = None
# Get loss function
loss_func = get_loss_func(args)
# Set up test set evaluation
test_smiles, test_targets = test_data.smiles(), test_data.targets()
if args.dataset_type == 'multiclass':
sum_test_preds = np.zeros((len(test_smiles), args.num_tasks, args.multiclass_num_classes))
else:
sum_test_preds = np.zeros((len(test_smiles), args.num_tasks))
# Automatically determine whether to cache
if len(data) <= args.cache_cutoff:
set_cache_graph(True)
num_workers = 0
else:
set_cache_graph(False)
num_workers = args.num_workers
# Create data loaders
train_data_loader = MoleculeDataLoader(
dataset=train_data,
batch_size=args.batch_size,
num_workers=num_workers,
class_balance=args.class_balance,
shuffle=True,
seed=args.seed
)
val_data_loader = MoleculeDataLoader(
dataset=val_data,
batch_size=args.batch_size,
num_workers=num_workers
)
test_data_loader = MoleculeDataLoader(
dataset=test_data,
batch_size=args.batch_size,
num_workers=num_workers
)
if args.class_balance:
debug(f'With class_balance, effective train size = {train_data_loader.iter_size:,}')
# Train ensemble of models
for model_idx in range(args.ensemble_size):
# Tensorboard writer
save_dir = os.path.join(args.save_dir, f'model_{model_idx}')
makedirs(save_dir)
try:
writer = SummaryWriter(log_dir=save_dir)
except:
writer = SummaryWriter(logdir=save_dir)
# Load/build model
if args.checkpoint_paths is not None:
debug(f'Loading model {model_idx} from {args.checkpoint_paths[model_idx]}')
model = load_checkpoint(args.checkpoint_paths[model_idx], logger=logger)
else:
debug(f'Building model {model_idx}')
model = MoleculeModel(args)
debug(model)
debug(f'Number of parameters = {param_count(model):,}')
if args.cuda:
debug('Moving model to cuda')
model = model.to(args.device)
# Ensure that model is saved in correct location for evaluation if 0 epochs
save_checkpoint(os.path.join(save_dir, MODEL_FILE_NAME), model, scaler,
features_scaler, atom_descriptor_scaler, bond_feature_scaler, args)
# Optimizers
optimizer = build_optimizer(model, args)
# Learning rate schedulers
scheduler = build_lr_scheduler(optimizer, args)
# Run training
best_score = float('inf') if args.minimize_score else -float('inf')
best_epoch, n_iter = 0, 0
for epoch in trange(args.epochs):
debug(f'Epoch {epoch}')
n_iter = train(
model=model,
data_loader=train_data_loader,
loss_func=loss_func,
optimizer=optimizer,
scheduler=scheduler,
args=args,
n_iter=n_iter,
logger=logger,
writer=writer
)
if isinstance(scheduler, ExponentialLR):
scheduler.step()
val_scores = evaluate(
model=model,
data_loader=val_data_loader,
num_tasks=args.num_tasks,
metrics=args.metrics,
dataset_type=args.dataset_type,
scaler=scaler,
logger=logger
)
for metric, scores in val_scores.items():
# Average validation score
avg_val_score = np.nanmean(scores)
debug(f'Validation {metric} = {avg_val_score:.6f}')
writer.add_scalar(f'validation_{metric}', avg_val_score, n_iter)
if args.show_individual_scores:
# Individual validation scores
for task_name, val_score in zip(args.task_names, scores):
debug(f'Validation {task_name} {metric} = {val_score:.6f}')
writer.add_scalar(f'validation_{task_name}_{metric}', val_score, n_iter)
# Save model checkpoint if improved validation score
avg_val_score = np.nanmean(val_scores[args.metric])
if args.minimize_score and avg_val_score < best_score or \
not args.minimize_score and avg_val_score > best_score:
best_score, best_epoch = avg_val_score, epoch
save_checkpoint(os.path.join(save_dir, MODEL_FILE_NAME), model, scaler, features_scaler,
atom_descriptor_scaler, bond_feature_scaler, args)
if args.split_sizes[1] == 0.0:
save_checkpoint(os.path.join(save_dir, MODEL_FILE_NAME), model,
scaler, features_scaler,
atom_descriptor_scaler, bond_feature_scaler, args)
# Evaluate on test set using model with best validation score
info(f'Model {model_idx} best validation {args.metric} = {best_score:.6f} on epoch {best_epoch}')
model = load_checkpoint(os.path.join(save_dir, MODEL_FILE_NAME), device=args.device, logger=logger)
test_preds = predict(
model=model,
data_loader=test_data_loader,
scaler=scaler
)
test_scores = evaluate_predictions(
preds=test_preds,
targets=test_targets,
num_tasks=args.num_tasks,
metrics=args.metrics,
dataset_type=args.dataset_type,
logger=logger
)
if len(test_preds) != 0:
sum_test_preds += np.array(test_preds)
# Average test score
for metric, scores in test_scores.items():
avg_test_score = np.nanmean(scores)
info(f'Model {model_idx} test {metric} = {avg_test_score:.6f}')
writer.add_scalar(f'test_{metric}', avg_test_score, 0)
if args.show_individual_scores:
# Individual test scores
for task_name, test_score in zip(args.task_names, scores):
info(f'Model {model_idx} test {task_name} {metric} = {test_score:.6f}')
writer.add_scalar(f'test_{task_name}_{metric}', test_score, n_iter)
writer.close()
# Evaluate ensemble on test set
avg_test_preds = (sum_test_preds / args.ensemble_size).tolist()
ensemble_scores = evaluate_predictions(
preds=avg_test_preds,
targets=test_targets,
num_tasks=args.num_tasks,
metrics=args.metrics,
dataset_type=args.dataset_type,
logger=logger
)
for metric, scores in ensemble_scores.items():
# Average ensemble score
avg_ensemble_test_score = np.nanmean(scores)
info(f'Ensemble test {metric} = {avg_ensemble_test_score:.6f}')
# Individual ensemble scores
if args.show_individual_scores:
for task_name, ensemble_score in zip(args.task_names, scores):
info(f'Ensemble test {task_name} {metric} = {ensemble_score:.6f}')
# Save scores
with open(os.path.join(args.save_dir, 'test_scores.json'), 'w') as f:
json.dump(ensemble_scores, f, indent=4, sort_keys=True)
# Optionally save test preds
if args.save_preds:
test_preds_dataframe = pd.DataFrame(data={'smiles': test_data.smiles()})
for i, task_name in enumerate(args.task_names):
test_preds_dataframe[task_name] = [pred[i] for pred in avg_test_preds]
test_preds_dataframe.to_csv(os.path.join(args.save_dir, 'test_preds.csv'), index=False)
return ensemble_scores