train

Apply OpenFF parameters to molecule, cluster conformers by RMSD and train.

Classes:

• TrainingFnArgs –

  Arguments for training functions.

• TrainFn –

  A protocol for training functions.

Functions:

• train_levenberg_marquardt –

  Iterate the training process using the Levenberg-Marquardt algorithm.

• train_adam –

  Iterate the training process using the Adam optimizer.

_TRAINING_FNS_REGISTRY module-attribute

_TRAINING_FNS_REGISTRY: dict[OptimiserName, TrainFn] = {}

Registry of training functions for different optimiser names.

TrainingFnArgs

Bases: TypedDict

Arguments for training functions.

TrainFn

Bases: Protocol

A protocol for training functions.

Methods:

• __call__ –

  Execute training function.

__call__

__call__(
    **kwargs: Unpack[TrainingFnArgs],
) -> tuple[Tensor, Trainable]

Execute training function.

Source code in presto/train.py

def __call__(
    self, **kwargs: Unpack[TrainingFnArgs]
) -> tuple[torch.Tensor, descent.train.Trainable]:
    """Execute training function."""

train_levenberg_marquardt

train_levenberg_marquardt(
    trainable_parameters: Tensor,
    initial_parameters: Tensor,
    trainable: Trainable,
    topologies: list[TensorTopology],
    datasets: list[Dataset],
    datasets_test: list[Dataset],
    settings: TrainingSettings,
    output_paths: dict[OutputType, PathLike],
    device: device,
) -> tuple[Tensor, Trainable]

Iterate the training process using the Levenberg-Marquardt algorithm.

Parameters:

• trainable_parameters (Tensor) –

  The parameters to be optimized.

• initial_parameters (Tensor) –

  The initial parameters before training.

• trainable (Trainable) –

  The trainable object containing the parameters.

• topologies (list[TensorTopology]) –

  The topologies of the systems.

• datasets (list[Dataset]) –

  The datasets to be used for training.

• datasets_test (list[Dataset]) –

  The datasets to be used for testing.

• settings (TrainingSettings) –

  The settings object containing training parameters.

• output_paths (dict[OutputType, PathLike]) –

  A mapping of output types to filesystem paths. The following keys are expected: OutputType.TENSORBOARD and OutputType.TRAINING_METRICS.

• device (device) –

  The device to perform training on.

Returns:

tuple[torch.Tensor, descent.train.Trainable] The updated parameters and the trainable object.

Source code in presto/train.py

@_register_training_fn("lm")
def train_levenberg_marquardt(
    trainable_parameters: torch.Tensor,
    initial_parameters: torch.Tensor,
    trainable: descent.train.Trainable,
    topologies: list[smee.TensorTopology],
    datasets: list[datasets.Dataset],
    datasets_test: list[datasets.Dataset],
    settings: TrainingSettings,
    output_paths: dict[OutputType, PathLike],
    device: torch.device,
) -> tuple[torch.Tensor, descent.train.Trainable]:
    """Iterate the training process using the Levenberg-Marquardt algorithm.

    Parameters
    ----------
    trainable_parameters : torch.Tensor
        The parameters to be optimized.
    initial_parameters : torch.Tensor
        The initial parameters before training.
    trainable : descent.train.Trainable
        The trainable object containing the parameters.
    topologies : list[smee.TensorTopology]
        The topologies of the systems.
    datasets : list[datasets.Dataset]
        The datasets to be used for training.
    datasets_test : list[datasets.Dataset]
        The datasets to be used for testing.
    settings : TrainingSettings
        The settings object containing training parameters.
    output_paths : dict[OutputType, PathLike]
        A mapping of output types to filesystem paths. The following keys are
        expected: ``OutputType.TENSORBOARD`` and ``OutputType.TRAINING_METRICS``.
    device : torch.device
        The device to perform training on.

    Returns:
    -------
    tuple[torch.Tensor, descent.train.Trainable]
        The updated parameters and the trainable object.
    """
    # Warn the user that LM needs more testing
    logger.warning(
        "Levenberg-Marquardt optimisation is an experimental feature and may not "
        "work as expected."
    )

    # Make sure we have all the required output paths and no others
    if set(output_paths.keys()) != settings.output_types:
        raise ValueError(
            f"Output paths must contain exactly the keys {settings.output_types}"
        )

    # Run the training with the LM optimiser
    lm_config = descent.optim.LevenbergMarquardtConfig(
        mode="adaptive", n_convergence_criteria=2, max_steps=settings.n_epochs
    )

    closure_fn = get_loss_closure_fn(
        datasets,
        trainable,
        trainable_parameters,
        initial_parameters,
        topologies,
        settings.regularisation_target,
    )

    correct_fn = trainable.clamp

    # Create report function that computes metrics consistently with train_adam
    report_fn = functools.partial(
        report,
        trainable=trainable,
        topologies=topologies,
        datasets_train=datasets,
        datasets_test=datasets_test,
        initial_parameters=initial_parameters,
        regularisation_target=settings.regularisation_target,
        metrics_file=output_paths[OutputType.TRAINING_METRICS],
        experiment_dir=Path(output_paths[OutputType.TENSORBOARD]),
    )

    trainable_parameters = descent.optim.levenberg_marquardt(
        trainable_parameters, lm_config, closure_fn, correct_fn, report_fn
    )
    trainable_parameters.requires_grad_(True)

    return trainable_parameters, trainable

train_adam

train_adam(
    trainable_parameters: Tensor,
    initial_parameters: Tensor,
    trainable: Trainable,
    topologies: list[TensorTopology],
    datasets: list[Dataset],
    datasets_test: list[Dataset],
    settings: TrainingSettings,
    output_paths: dict[OutputType, PathLike],
    device: device,
) -> tuple[Tensor, Trainable]

Iterate the training process using the Adam optimizer.

Parameters:

• trainable_parameters (Tensor) –

  The parameters to be optimized.

• initial_parameters (Tensor) –

  The initial parameters before training.

• trainable (Trainable) –

  The trainable object containing the parameters.

• topologies (list[TensorTopology]) –

  The topologies of the systems.

• datasets (list[Dataset]) –

  The datasets to be used for training.

• datasets_test (list[Dataset]) –

  The datasets to be used for testing.

• settings (TrainingSettings) –

  The settings object containing training parameters.

• output_paths (dict[OutputType, PathLike]) –

  A mapping of output types to filesystem paths. The following keys are expected: OutputType.TENSORBOARD and OutputType.TRAINING_METRICS.

• device (device) –

  The device to perform training on.

Returns:

tuple[torch.Tensor, descent.train.Trainable] The updated parameters and the trainable object.

Source code in presto/train.py

@_register_training_fn("adam")
def train_adam(
    trainable_parameters: torch.Tensor,
    initial_parameters: torch.Tensor,
    trainable: descent.train.Trainable,
    topologies: list[smee.TensorTopology],
    datasets: list[datasets.Dataset],
    datasets_test: list[datasets.Dataset],
    settings: TrainingSettings,
    output_paths: dict[OutputType, PathLike],
    device: torch.device,
) -> tuple[torch.Tensor, descent.train.Trainable]:
    """Iterate the training process using the Adam optimizer.

    Parameters
    ----------
    trainable_parameters : torch.Tensor
        The parameters to be optimized.
    initial_parameters : torch.Tensor
        The initial parameters before training.
    trainable : descent.train.Trainable
        The trainable object containing the parameters.
    topologies : list[smee.TensorTopology]
        The topologies of the systems.
    datasets : list[datasets.Dataset]
        The datasets to be used for training.
    datasets_test : list[datasets.Dataset]
        The datasets to be used for testing.
    settings : TrainingSettings
        The settings object containing training parameters.
    output_paths : dict[OutputType, PathLike]
        A mapping of output types to filesystem paths. The following keys are
        expected: ``OutputType.TENSORBOARD`` and ``OutputType.TRAINING_METRICS``.
    device : torch.device
        The device to perform training on.

    Returns:
    -------
    tuple[torch.Tensor, descent.train.Trainable]
        The updated parameters and the trainable object.
    """
    # Make sure we have all the required output paths and no others
    if set(output_paths.keys()) != settings.output_types:
        raise ValueError(
            f"Output paths must contain exactly the keys {settings.output_types}"
        )

    # run the ML training

    with open(output_paths[OutputType.TRAINING_METRICS], "w") as metrics_file:
        with open_writer(Path(output_paths[OutputType.TENSORBOARD])) as writer:
            optimizer = torch.optim.Adam(
                [trainable_parameters], lr=settings.learning_rate, amsgrad=True
            )
            scheduler = torch.optim.lr_scheduler.ExponentialLR(
                optimizer, gamma=settings.learning_rate_decay
            )
            for v in tensorboardX.writer.hparams(
                {"optimizer": "Adam", "lr": settings.learning_rate}, {}
            ):
                writer.file_writer.add_summary(v)
            progress = Progress(
                TextColumn("[progress.description]{task.description}"),
                BarColumn(),
                MofNCompleteColumn(),
                TimeRemainingColumn(),
            )
            with progress:
                for i in progress.track(
                    range(settings.n_epochs),
                    description="Optimising MM parameters",
                ):
                    # Memory optimisation: Use manual gradient computation which
                    # computes gradients per-molecule and detaches immediately
                    losses_train, gradient = compute_overall_loss_and_grad(
                        datasets,
                        trainable,
                        trainable_parameters,
                        initial_parameters,
                        topologies,
                        settings.regularisation_target,
                        device,
                    )

                    logger.debug(
                        f"Epoch {i}: Training Weighted Loss: "
                        f"Energy={losses_train.energy.item():.4f} "
                        f"Forces={losses_train.forces.item():.4f} "
                        f"Reg={losses_train.regularisation.item():.4f}"
                    )
                    if i % 10 == 0:
                        losses_test, _ = compute_overall_loss_and_grad(
                            datasets_test,
                            trainable,
                            trainable_parameters,
                            initial_parameters,
                            topologies,
                            settings.regularisation_target,
                            device,
                            compute_grad=False,
                        )

                        write_metrics(
                            i,
                            losses_train,
                            losses_test,
                            writer,
                            metrics_file,
                        )

                    trainable_parameters.grad = gradient
                    optimizer.step()
                    optimizer.zero_grad(set_to_none=True)
                    trainable.clamp(trainable_parameters)

                    if i % settings.learning_rate_decay_step == 0:
                        scheduler.step()

        # Required to avoid filling up the GPU memory between iterations
        # TODO: Find a better way to do this.
        if torch.cuda.is_available():
            torch.cuda.synchronize()
            torch.cuda.empty_cache()

        # some book-keeping and outputting
        losses_train, _ = compute_overall_loss_and_grad(
            datasets,
            trainable,
            trainable_parameters,
            initial_parameters,
            topologies,
            settings.regularisation_target,
            device,
            compute_grad=False,
        )
        losses_test, _ = compute_overall_loss_and_grad(
            datasets_test,
            trainable,
            trainable_parameters,
            initial_parameters,
            topologies,
            settings.regularisation_target,
            device,
            compute_grad=False,
        )

        write_metrics(
            settings.n_epochs, losses_train, losses_test, writer, metrics_file
        )

        return trainable_parameters, trainable