Metrics

from lightning import pytorch as pl
import numpy as np
from numpy.typing import ArrayLike
import pandas as pd
from pathlib import Path
import torch
from torch import Tensor
import torchmetrics
import logging

from chemprop import data, models, nn
from chemprop.nn.metrics import ChempropMetric, MetricRegistry

logger = logging.getLogger(__name__)

Available metric functions

Chemprop provides several metrics. The functions calculate a single value that serves as a measure of model performance. Users only need to select the metric(s) to use. The rest of the details are handled by Chemprop and the lightning trainer, which logs all metric values to the trainer logger (defaults to TensorBoard) for the validation and test sets. Note that the validation metrics are in the scaled space while the test metrics are in the original target space.

See also loss functions which are the same as metrics, except used to optimize the model and therefore required to be differentiable.

for metric in MetricRegistry:
    print(metric)

mse
mae
rmse
bounded-mse
bounded-mae
bounded-rmse
r2
binary-mcc
multiclass-mcc
roc
prc
accuracy
f1

Specifying metrics

Each FFN predictor has a default metric. If you want different metrics reported, you can give a list of metrics to the model at creation. Note that the list of metrics is used in place of the default metric and not in addition to the default metric.

from chemprop.nn.metrics import MSE, MAE, RMSE

metrics = [MSE(), MAE(), RMSE()]
model = models.MPNN(nn.BondMessagePassing(), nn.MeanAggregation(), nn.RegressionFFN(), metrics=metrics)

Accumulating metrics

Chemprop metrics are based on Metric from torchmetrics which stores the information from each batch that is needed to calculate the metric over the whole validation or test set.

smis = ["C" * i for i in range(1, 11)]
ys = np.random.rand(len(smis), 1)
dset = data.MoleculeDataset([data.MoleculeDatapoint.from_smi(smi, y) for smi, y in zip(smis, ys)])
dataloader = data.build_dataloader(dset, shuffle=False, batch_size=2)

trainer = pl.Trainer(logger=False, enable_checkpointing=False, max_epochs=1)
result_when_batched = trainer.test(model, dataloader)
preds = trainer.predict(model, dataloader)
preds = torch.concat(preds)

GPU available: False, used: False
TPU available: False, using: 0 TPU cores

/home/knathan/anaconda3/envs/chemprop/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:434: The 'test_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃        Test metric        ┃       DataLoader 0        ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│         test/mae          │    0.6082264184951782     │
│         test/mse          │    0.41469845175743103    │
│         test/rmse         │    0.6439708471298218     │
└───────────────────────────┴───────────────────────────┘

/home/knathan/anaconda3/envs/chemprop/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:434: The 'predict_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

result_when_not_batched = RMSE()(preds, torch.from_numpy(dset.Y), None, None, None, None)
print("Batch / Not Batched")
print(f"{result_when_batched[0]['test/rmse']:.4f} / {result_when_not_batched.item():.4f}")

Batch / Not Batched
0.6440 / 0.6440

Batch normalization

It is worth noting that if your model has a batch normalization layer, the computed metric will be different depending on if the model is in training or evaluation mode. When a batch normalization layer is training, it uses a biased estimator to calculate the standard deviation, but the value stored and used during evaluation is calculated with the unbiased estimator. Lightning takes care of this if the Trainer() is used.

Regression

There are several metric options for regression. MSE is the default. There are also bounded versions (except for r2), similar to the bounded versions of the loss functions.

from chemprop.nn.metrics import MSE, MAE, RMSE, R2Score

from chemprop.nn.metrics import BoundedMAE, BoundedMSE, BoundedRMSE

Classification

There are metrics for both binary and multiclass classification.

from chemprop.nn.metrics import (
    BinaryAUROC,
    BinaryAUPRC,
    BinaryAccuracy,
    BinaryF1Score,
    BinaryMCCMetric,
)

from chemprop.nn.metrics import MulticlassMCCMetric

Spectra

Spectral information divergence and wasserstein (earthmovers distance) are often used for spectral predictions

from chemprop.nn.metrics import SID, Wasserstein

Custom metrics

Chemprop metrics are instances of chemprop.nn.metrics.ChempropMetric, which inherits from torchmetrics.Metric. Custom loss functions need to follow the interface of both ChempropMetric and Metric. Start with a Metric either by importing an existing one from torchmetrics or by creating your own by following the instructions on the torchmetrics website. Then make the following changes:

1. Allow for task weights to be passed to the __init__ method.

2. Allow for the update method to be given preds, targets, mask, weights, lt_mask, gt_mask in that order.

3. Provide an alias property, which is used to identify the metric value in the logs.

• preds: A Tensor of the model’s predictions with dimension 0 being the batch dimension and dimension 1 being the task dimension.

• targets: A Tensor of the target values with dimension 0 being the batch dimension and dimension 1 being the task dimension.

• mask: A Tensor of the same shape as targets with Trues where the target value is present and finite and False where it is not.

• weights: Usually ignored in metrics.

• lt_mask: A Tensor of the same shape as targets with Trues where the target value is a “less than” target value and False where it is not.

• gt_mask: A Tensor of the same shape as targets with Trues where the target value is a “greater than” target value and False where it is not.

class ChempropMulticlassAUROC(torchmetrics.classification.MulticlassAUROC):
    def __init__(self, task_weights: ArrayLike = 1.0, **kwargs):
        super().__init__(**kwargs)
        self.task_weights = torch.as_tensor(task_weights, dtype=torch.float).view(1, -1)
        if (self.task_weights != 1.0).any():
            logger.warn("task_weights were provided but are ignored by metric "
                          f"{self.__class__.__name__}. Got {task_weights}")

    def update(self, preds: Tensor, targets: Tensor, mask: Tensor | None = None, *args, **kwargs):
        if mask is None:
            mask = torch.ones_like(targets, dtype=torch.bool)

        super().update(preds[mask], targets[mask].long())

    @property
    def alias(self) -> str:
        return "multiclass_auroc"

Alternatively, if your metric can return a value for every task for every data point (i.e. not reduced in the task or batch dimension), you can inherit from chemprop.nn.metrics.ChempropMetric and just override the _calc_unreduced_loss method (and if needed the __init__ method).

class BoundedNormalizedMSEPlus1(ChempropMetric):
    def __init__(self, task_weights = None, norm: float = 1.0):
        super().__init__(task_weights)
        norm = torch.as_tensor(norm)
        self.register_buffer("norm", norm)

    def _calc_unreduced_loss(self, preds, targets, mask, weights, lt_mask, gt_mask) -> Tensor:
        preds = torch.where((preds < targets) & lt_mask, targets, preds)
        preds = torch.where((preds > targets) & gt_mask, targets, preds)

        return torch.sum((preds - targets) ** 2) / self.norm + 1

Example

Set up the dataset

chemprop_dir = Path.cwd().parents[3]
input_path = chemprop_dir / "tests" / "data" / "classification" / "mol_multiclass.csv"
df_input = pd.read_csv(input_path)
smis = df_input.loc[:, "smiles"].values
ys = df_input.loc[:, ["activity"]].values
all_data = [data.MoleculeDatapoint.from_smi(smi, y) for smi, y in zip(smis, ys)]
train_indices, val_indices, test_indices = data.make_split_indices(all_data, "random", (0.8, 0.1, 0.1))
train_data, val_data, test_data = data.split_data_by_indices(
    all_data, train_indices, val_indices, test_indices
)
train_dset = data.MoleculeDataset(train_data[0])
val_dset = data.MoleculeDataset(val_data[0])
test_dset = data.MoleculeDataset(test_data[0])
train_loader = data.build_dataloader(train_dset)
val_loader = data.build_dataloader(val_dset, shuffle=False)
test_loader = data.build_dataloader(test_dset, shuffle=False)

The return type of make_split_indices has changed in v2.1 - see help(make_split_indices)

Make a model with a custom loss function

n_classes = max(ys).item() + 1

metrics = [ChempropMulticlassAUROC(num_classes = n_classes)]

model = models.MPNN(
    nn.BondMessagePassing(),
    nn.NormAggregation(),
    nn.MulticlassClassificationFFN(n_classes = n_classes),
    metrics = metrics
    )

Run training

trainer = pl.Trainer(max_epochs=2)
trainer.fit(model, train_loader, val_loader)
trainer.test(model, test_loader, weights_only=False)  # weights_only=False is only required with pytorch lightning version 2.6.0 or newer

💡 Tip: For seamless cloud uploads and versioning, try installing [litmodels](https://pypi.org/project/litmodels/) to enable LitModelCheckpoint, which syncs automatically with the Lightning model registry.
GPU available: False, used: False
TPU available: False, using: 0 TPU cores
Loading `train_dataloader` to estimate number of stepping batches.
/home/knathan/anaconda3/envs/chemprop/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.
/home/knathan/anaconda3/envs/chemprop/lib/python3.11/site-packages/lightning/pytorch/loops/fit_loop.py:317: The number of training batches (7) is smaller than the logging interval Trainer(log_every_n_steps=50). Set a lower value for log_every_n_steps if you want to see logs for the training epoch.

┏━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━┳━━━━━━━┳━━━━━━━┓
┃   ┃ Name            ┃ Type                        ┃ Params ┃ Mode  ┃ FLOPs ┃
┡━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━╇━━━━━━━╇━━━━━━━┩
│ 0 │ message_passing │ BondMessagePassing          │  227 K │ train │     0 │
│ 1 │ agg             │ NormAggregation             │      0 │ train │     0 │
│ 2 │ bn              │ Identity                    │      0 │ train │     0 │
│ 3 │ predictor       │ MulticlassClassificationFFN │ 91.2 K │ train │     0 │
│ 4 │ X_d_transform   │ Identity                    │      0 │ train │     0 │
│ 5 │ metrics         │ ModuleList                  │      0 │ train │     0 │
└───┴─────────────────┴─────────────────────────────┴────────┴───────┴───────┘

Trainable params: 318 K
Non-trainable params: 0
Total params: 318 K
Total estimated model params size (MB): 1
Modules in train mode: 24
Modules in eval mode: 0
Total FLOPs: 0

/home/knathan/anaconda3/envs/chemprop/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/data_connec
tor.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the
value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

/home/knathan/anaconda3/envs/chemprop/lib/python3.11/site-packages/lightning/pytorch/core/saving.py:365: Skipping 'metrics' parameter because it is not possible to safely dump to YAML.

/home/knathan/anaconda3/envs/chemprop/lib/python3.11/site-packages/torchmetrics/utilities/prints.py:43:
UserWarning: No positive samples in targets, true positive value should be meaningless. Returning zero tensor in
true positive score
  warnings.warn(*args, **kwargs)

`Trainer.fit` stopped: `max_epochs=2` reached.

┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃        Test metric        ┃       DataLoader 0        ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│   test/multiclass_auroc   │    0.6266666650772095     │
└───────────────────────────┴───────────────────────────┘

/home/knathan/anaconda3/envs/chemprop/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:434: The 'test_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

[{'test/multiclass_auroc': 0.6266666650772095}]