Loss functions

Loss functions#

[1]:

import warnings

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

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

Available functions#

Chemprop provides several loss functions. The derivatives of these differentiable functions are used to update the model weights. Users only need to select the loss function to use. The rest of the details are handled by Chemprop and the lightning trainer, which reports the training and validation loss during model fitting.

See also metrics which are the same as loss functions, but potentially non-differentiable and used to measure the performance of a model.

[2]:

for lossfunction in LossFunctionRegistry:
    print(lossfunction)

mse
mae
rmse
bounded-mse
bounded-mae
bounded-rmse
mve
evidential
bce
ce
binary-mcc
multiclass-mcc
dirichlet
sid
earthmovers
wasserstein
quantile
pinball
quantile-point
pinball-point

Task weights#

A model can make predictions of multiple targets/tasks at the same time. For example, a model may predict both solubility and melting point. Task weights can be specified when some of the tasks are more important to get accurate than others. The weight for each task defaults to 1.

[3]:

from chemprop.nn.metrics import MSE

predictor = nn.RegressionFFN(criterion=MSE(task_weights=[0.1, 0.5, 1.0]))
model = models.MPNN(nn.BondMessagePassing(), nn.MeanAggregation(), predictor)
predictor.criterion

[3]:

MSE(task_weights=[[0.10000000149011612, 0.5, 1.0]])

Mean squared error and bounded mean square error#

MSE is the default loss function for regression tasks.

[4]:

predictor = nn.RegressionFFN()
model = models.MPNN(nn.BondMessagePassing(), nn.MeanAggregation(), predictor)
predictor.criterion

[4]:

MSE(task_weights=[[1.0]])

BoundedMSE is useful when the target values have “less than” or “greater than” behavior, e.g. the prediction is correct as long as it is below/above a target value. Datapoints have a less than/greater than property that keeps track of bounded targets. Note that, like target values, the less than and greater than masks used to make datapoints are 1-D numpy arrays of bools instead of a single bool. This is because a single datapoint can have multiple target values and the less than/greater than masks are defined for each target value separately.

[5]:

from chemprop.nn.metrics import BoundedMSE

smis = ["C" * i for i in range(1, 6)]
ys = np.random.rand(len(smis), 1)
lt_mask = np.array([[True], [False], [False], [False], [True]])
gt_mask = np.array([[False], [True], [False], [True], [False]])
datapoints = [
    data.MoleculeDatapoint.from_smi(smi, y, lt_mask=lt, gt_mask=gt)
    for smi, y, lt, gt in zip(smis, ys, lt_mask, gt_mask)
]
bounded_dataset = data.MoleculeDataset(datapoints)
bounded_dataset.lt_mask

[5]:

array([[ True],
       [False],
       [False],
       [False],
       [ True]])

[6]:

predictor = nn.RegressionFFN(criterion=BoundedMSE())
model = models.MPNN(nn.BondMessagePassing(), nn.MeanAggregation(), predictor)

Binary cross entropy and cross entropy#

BCELoss is the default loss function for binary classification and CrossEntropyLoss is the default for multiclass classification.

[7]:

predictor = nn.BinaryClassificationFFN()
predictor.criterion

[7]:

BCELoss(task_weights=[[1.0]])

[8]:

predictor = nn.MulticlassClassificationFFN(n_classes=3)
predictor.criterion

[8]:

CrossEntropyLoss(task_weights=[[1.0]])

Matthews correlation coefficient#

MCC loss is useful for imbalanced classification data. An optimal MCC is 1, so the loss function version of MCC returns 1 - MCC.

[9]:

from chemprop.nn.metrics import BinaryMCCLoss, MulticlassMCCLoss

Uncertainty#

Various methods for estimating uncertainty in predictions are available. These methods often use specific loss functions.

[10]:

from chemprop.nn.metrics import MVELoss, EvidentialLoss, DirichletLoss, QuantileLoss

Spectral loss functions#

Spectral information divergence and wasserstein (earthmover’s distance) are often used for spectral predictions.

[11]:

from chemprop.nn.metrics import SID, Wasserstein

Custom loss functions#

Chemprop loss functions 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.

  • preds: A Tensor of the model’s predictions with dimension 0 being the batch dimension and dimension 1 being the task dimension. Dimension 2 exists for uncertainty estimation or multiclass predictions and is either used for uncertainty parameters or multiclass logits.

  • 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: A Tensor of the weights for each data point in the loss function. This is useful when some data points are more important than others.

  • 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.

[12]:

class ChempropMulticlassHingeLoss(torchmetrics.classification.MulticlassHingeLoss):
    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():
            warnings.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())

Alternatively, if your loss function 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).

[13]:

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

[14]:

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

[15]:

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

loss_function = ChempropMulticlassHingeLoss(num_classes = n_classes)
ffn = nn.MulticlassClassificationFFN(n_classes = n_classes, criterion = loss_function)

model = models.MPNN(nn.BondMessagePassing(), nn.NormAggregation(), ffn)

Run training

[16]:

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

/home/jackson/miniconda3/envs/chemprop_dev/lib/python3.12/site-packages/torch/cuda/__init__.py:734: UserWarning: Can't initialize NVML
  warnings.warn("Can't initialize NVML")
💡 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/jackson/miniconda3/envs/chemprop_dev/lib/python3.12/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=15` in the `DataLoader` to improve performance.
/home/jackson/miniconda3/envs/chemprop_dev/lib/python3.12/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/jackson/miniconda3/envs/chemprop_dev/lib/python3.12/site-packages/rich/live.py:231: UserWarning: install
"ipywidgets" for Jupyter support
  warnings.warn('install "ipywidgets" for Jupyter support')

/home/jackson/miniconda3/envs/chemprop_dev/lib/python3.12/site-packages/lightning/pytorch/trainer/connectors/data_c
onnector.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=15` in the `DataLoader` to improve performance.

/home/jackson/miniconda3/envs/chemprop_dev/lib/python3.12/site-packages/torch/cuda/__init__.py:734: UserWarning: Can't initialize NVML
  warnings.warn("Can't initialize NVML")
`Trainer.fit` stopped: `max_epochs=2` reached.













/home/jackson/miniconda3/envs/chemprop_dev/lib/python3.12/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=15` in the `DataLoader` to improve performance.













┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃        Test metric        ┃       DataLoader 0        ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│    test/multiclass-mcc    │            0.0            │
└───────────────────────────┴───────────────────────────┘































[16]:







[{'test/multiclass-mcc': 0.0}]