Training and Predicting

chemprop.train contains functions to train and make predictions with message passing neural networks.

Train

chemprop.train.train.py trains a model for a single epoch.

chemprop.train.train.train(model: MoleculeModel, data_loader: MoleculeDataLoader, loss_func: Callable, optimizer: Optimizer, scheduler: _LRScheduler, args: TrainArgs, n_iter: int = 0, atom_bond_scaler: AtomBondScaler | None = None, logger: Logger | None = None, writer: SummaryWriter | None = None) → int

Trains a model for an epoch.

Parameters:

• model – A MoleculeModel.

• data_loader – A MoleculeDataLoader.

• loss_func – Loss function.

• optimizer – An optimizer.

• scheduler – A learning rate scheduler.

• args – A TrainArgs object containing arguments for training the model.

• n_iter – The number of iterations (training examples) trained on so far.

• atom_bond_scaler – A AtomBondScaler fitted on the atomic/bond targets.

• logger – A logger for recording output.

• writer – A tensorboardX SummaryWriter.

Returns:

The total number of iterations (training examples) trained on so far.

Run Training

chemprop.train.run_training.py loads data, initializes the model, and runs training, validation, and testing of the model.

chemprop.train.run_training.run_training(args: TrainArgs, data: MoleculeDataset, logger: Logger | None = None) → Dict[str, List[float]]

Loads data, trains a Chemprop model, and returns test scores for the model checkpoint with the highest validation score.

Parameters:

• args – A TrainArgs object containing arguments for loading data and training the Chemprop model.

• data – A MoleculeDataset containing the data.

• logger – A logger to record output.

Returns:

A dictionary mapping each metric in args.metrics to a list of values for each task.

Cross-Validation

chemprop.train.cross_validate.py provides an outer loop around chemprop.train.run_training.py that runs training and evaluating for each of several splits of the data.

chemprop.train.cross_validate.chemprop_train() → None

Parses Chemprop training arguments and trains (cross-validates) a Chemprop model.

This is the entry point for the command line command chemprop_train.

chemprop.train.cross_validate.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.

Parameters:

• args – A TrainArgs object containing arguments for loading data and training the Chemprop model.

• train_func – Function which runs training.

Returns:

A tuple containing the mean and standard deviation performance across folds.

Predict

chemprop.train.predict.py uses a trained model to make predicts on data.

chemprop.train.predict.predict(model: MoleculeModel, data_loader: MoleculeDataLoader, disable_progress_bar: bool = False, scaler: StandardScaler | None = None, atom_bond_scaler: AtomBondScaler | None = None, return_unc_parameters: bool = False, dropout_prob: float = 0.0) → List[List[float]]

Makes predictions on a dataset using an ensemble of models.

Parameters:

• model – A MoleculeModel.

• data_loader – A MoleculeDataLoader.

• disable_progress_bar – Whether to disable the progress bar.

• scaler – A StandardScaler object fit on the training targets.

• atom_bond_scaler – A AtomBondScaler fitted on the atomic/bond targets.

• return_unc_parameters – A bool indicating whether additional uncertainty parameters would be returned alongside the mean predictions.

• dropout_prob – For use during uncertainty prediction only. The propout probability used in generating a dropout ensemble.

Returns:

A list of lists of predictions. The outer list is molecules while the inner list is tasks. If returning uncertainty parameters as well, it is a tuple of lists of lists, of a length depending on how many uncertainty parameters are appropriate for the loss function.

Make Predictions

chemprop.train.make_predictions.py is a wrapper aoround chemprop.train.predict.py which loads data, loads a trained model, makes predictions, and saves those predictions.

chemprop.train.make_predictions.chemprop_predict() → None

Parses Chemprop predicting arguments and runs prediction using a trained Chemprop model.

This is the entry point for the command line command chemprop_predict.

chemprop.train.make_predictions.load_data(args: PredictArgs, smiles: List[List[str]])

Function to load data from a list of smiles or a file.

Parameters:

• args – A PredictArgs object containing arguments for loading data and a model and making predictions.

• smiles – A list of list of smiles, or None if data is to be read from file

Returns:

A tuple of a MoleculeDataset containing all datapoints, a MoleculeDataset containing only valid datapoints, a MoleculeDataLoader and a dictionary mapping full to valid indices.

chemprop.train.make_predictions.load_model(args: PredictArgs, generator: bool = False)

Function to load a model or ensemble of models from file. If generator is True, a generator of the respective model and scaler objects is returned (memory efficient), else the full list (holding all models in memory, necessary for preloading).

Parameters:

• args – A PredictArgs object containing arguments for loading data and a model and making predictions.

• generator – A boolean to return a generator instead of a list of models and scalers.

Returns:

A tuple of updated prediction arguments, training arguments, a list or generator object of models, a list or generator object of scalers, the number of tasks and their respective names.

chemprop.train.make_predictions.make_predictions(args: PredictArgs, smiles: List[List[str]] = None, model_objects: Tuple[PredictArgs, TrainArgs, List[MoleculeModel], List[StandardScaler | AtomBondScaler], int, List[str]] = None, calibrator: UncertaintyCalibrator = None, return_invalid_smiles: bool = True, return_index_dict: bool = False, return_uncertainty: bool = False) → List[List[float | None]]

Loads data and a trained model and uses the model to make predictions on the data.

If SMILES are provided, then makes predictions on smiles. Otherwise makes predictions on args.test_data.

Parameters:

• args – A PredictArgs object containing arguments for loading data and a model and making predictions.

• smiles – List of list of SMILES to make predictions on.

• model_objects – Tuple of output of load_model function which can be called separately outside this function. Preloaded model objects should have used the non-generator option for load_model if the objects are to be used multiple times or are intended to be used for calibration as well.

• calibrator – A :class: ~chemprop.uncertainty.UncertaintyCalibrator object, for use in calibrating uncertainty predictions. Can be preloaded and provided as a function input or constructed within the function from arguments. The models and scalers used to initiate the calibrator must be lists instead of generators if the same calibrator is to be used multiple times or if the same models and scalers objects are also part of the provided model_objects input.

• return_invalid_smiles – Whether to return predictions of “Invalid SMILES” for invalid SMILES, otherwise will skip them in returned predictions.

• return_index_dict – Whether to return the prediction results as a dictionary keyed from the initial data indexes.

• return_uncertainty – Whether to return uncertainty predictions alongside the model value predictions.

Returns:

A list of lists of target predictions. If returning uncertainty, a tuple containing first prediction values then uncertainty estimates.

chemprop.train.make_predictions.predict_and_save(args: PredictArgs, train_args: TrainArgs, test_data: MoleculeDataset, task_names: List[str], num_tasks: int, test_data_loader: MoleculeDataLoader, full_data: MoleculeDataset, full_to_valid_indices: dict, models: List[MoleculeModel], scalers: List[StandardScaler | AtomBondScaler], num_models: int, calibrator: UncertaintyCalibrator | None = None, return_invalid_smiles: bool = False, save_results: bool = True)

Function to predict with a model and save the predictions to file.

Parameters:

• args – A PredictArgs object containing arguments for loading data and a model and making predictions.

• train_args – A TrainArgs object containing arguments for training the model.

• test_data – A MoleculeDataset containing valid datapoints.

• task_names – A list of task names.

• num_tasks – Number of tasks.

• test_data_loader – A MoleculeDataLoader to load the test data.

• full_data – A MoleculeDataset containing all (valid and invalid) datapoints.

• full_to_valid_indices – A dictionary dictionary mapping full to valid indices.

• models – A list or generator object of MoleculeModels.

• scalers – A list or generator object of StandardScaler objects.

• num_models – The number of models included in the models and scalers input.

• calibrator – A :class: ~chemprop.uncertainty.UncertaintyCalibrator object, for use in calibrating uncertainty predictions.

• return_invalid_smiles – Whether to return predictions of “Invalid SMILES” for invalid SMILES, otherwise will skip them in returned predictions.

• save_results – Whether to save the predictions in a csv. Function returns the predictions regardless.

Returns:

A list of lists of target predictions.

chemprop.train.make_predictions.set_features(args: PredictArgs, train_args: TrainArgs)

Function to set extra options.

Parameters:

• args – A PredictArgs object containing arguments for loading data and a model and making predictions.

• train_args – A TrainArgs object containing arguments for training the model.

Evaluate

chemprop.train.evaluate.py contains functions for evaluating the quality of predictions by comparing them to the true values.

chemprop.train.evaluate.evaluate(model: MoleculeModel, data_loader: MoleculeDataLoader, num_tasks: int, metrics: List[str], dataset_type: str, scaler: StandardScaler | None = None, quantiles: List[float] | None = None, atom_bond_scaler: AtomBondScaler | None = None, logger: Logger | None = None) → Dict[str, List[float]]

Evaluates an ensemble of models on a dataset by making predictions and then evaluating the predictions.

Parameters:

• model – A MoleculeModel.

• data_loader – A MoleculeDataLoader.

• num_tasks – Number of tasks.

• metrics – A list of names of metric functions.

• dataset_type – Dataset type.

• scaler – A StandardScaler object fit on the training targets.

• quantiles – A list of quantiles for use in pinball evaluation of quantile_interval metric.

• atom_bond_scaler – A AtomBondScaler fitted on the atomic/bond targets.

• logger – A logger to record output.

Returns:

A dictionary mapping each metric in metrics to a list of values for each task.

chemprop.train.evaluate.evaluate_predictions(preds: List[List[float]], targets: List[List[float]], num_tasks: int, metrics: List[str], dataset_type: str, is_atom_bond_targets: bool = False, gt_targets: List[List[bool]] | None = None, lt_targets: List[List[bool]] | None = None, quantiles: List[float] | None = None, logger: Logger | None = None) → Dict[str, List[float]]

Evaluates predictions using a metric function after filtering out invalid targets.

Parameters:

• preds – A list of lists of shape (data_size, num_tasks) with model predictions.

• targets – A list of lists of shape (data_size, num_tasks) with targets.

• num_tasks – Number of tasks.

• metrics – A list of names of metric functions.

• dataset_type – Dataset type.

• is_atom_bond_targets – Boolean whether this is atomic/bond properties prediction.

• gt_targets – A list of lists of booleans indicating whether the target is an inequality rather than a single value.

• lt_targets – A list of lists of booleans indicating whether the target is an inequality rather than a single value.

• quantiles – A list of quantiles for use in pinball evaluation of quantile_interval metric.

• logger – A logger to record output.

Returns:

A dictionary mapping each metric in metrics to a list of values for each task.