Callbacks

from pathlib import Path

from lightning.pytorch.callbacks import Callback
import pandas as pd
import torch

from chemprop.callbacks import CallbackRegistry

Available callbacks

Chemprop uses PyTorch Lightning for training and prediction, which allows for the use of callbacks to add custom behavior during these processes. Callbacks can be used for a variety of purposes, such as logging, model interpretation, or saving additional artifacts.

Currently command line usage of callbacks is only supported for the predict subcommand (only one callback at a time). Here are the callbacks currently available in Chemprop:

for callback in CallbackRegistry:
    print(callback)

myerson

Custom callbacks

You can create your own custom callbacks to perform specific actions during the predict subcommands. A custom callback must inherit from lightning.pytorch.callbacks.Callback and can be registered with Chemprop’s CallbackRegistry to be accessible from the command line.

If you don’t want to change chemprop source code to register the custom callback, you can call the pl.Trainer manually with the callbacks=[ExampleCallback, ...] argument.

A callback has access to the pl.Trainer and pl.LightningModule objects, and can implement any of the hooks provided by PyTorch Lightning. For a full list of available hooks, please refer to the PyTorch Lightning documentation on callbacks.

Here is the basic structure of a Chemprop callback:

@CallbackRegistry.register("example_callback")
class ExampleCallback(Callback):
    def __init__(self, model_paths: list[Path], output: Path):
        super().__init__()
        print("ExampleCallback initialized!")
        print(f"  model_paths: {model_paths}")
        print(f"  output: {output}")

    def on_predict_epoch_start(self, trainer, pl_module):
        print("\nPrediction epoch started.")

    def on_predict_epoch_end(self, trainer, pl_module):
        print("Prediction epoch ended.")

    def on_predict_batch_start(self, trainer, pl_module, batch, batch_idx, dataloader_idx=0):
        print(f"  Processing batch {batch_idx}...")

    def on_predict_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0):
        print(f"  ...finished batch {batch_idx}.")

    def on_predict_start(self, trainer, pl_module):
        print("\nStarting prediction...")

    def on_predict_end(self, trainer, pl_module):
        print("...prediction finished.")

Key components:

1. Inheritance: Your callback class must inherit from lightning.pytorch.callbacks.Callback.

2. Registration: The @CallbackRegistry.register("your_callback_name") decorator makes your callback available via the --callback command-line argument.

3. ``__init__`` method: The constructor for callbacks used with the predict command receives model_paths and output as arguments, which are the paths to the model(s) being used and the specified output file for predictions. You can also accept additional parameters, which can be passed from the command line as a JSON string to --callback-params.

4. Hooks: Implement any of the on_* methods (hooks) from PyTorch Lightning to execute your code at different stages of the process. The example above shows several hooks available during prediction.

For a more complex, real-world example, you can look at chemprop.callbacks.interpret_callbacks.MyersonExplainerCallback, which calculates and saves atom-level model explanations during prediction.