chemprop.nn.ffn

Classes

FFN	A FFN is a differentiable function
MLP	An MLP is an FFN that implements the following function:
ConstrainerFFN	A ConstrainerFFN adjusts atom or bond property predictions to satisfy molecular

Module Contents

class chemprop.nn.ffn.FFN(*args, **kwargs)

Bases: torch.nn.Module

A FFN is a differentiable function \(f_\theta : \mathbb R^i \mapsto \mathbb R^o\)

Parameters:

• args (Any)

• kwargs (Any)

input_dim: int

output_dim: int

abstractmethod forward(X)

Parameters:

X (torch.Tensor)

Return type:

torch.Tensor

class chemprop.nn.ffn.MLP(*args: torch.nn.modules.module.Module)

class chemprop.nn.ffn.MLP(arg)

Bases: torch.nn.Sequential, FFN

An MLP is an FFN that implements the following function:

\[\begin{split}\mathbf h_0 &= \mathbf W_0 \mathbf x \,+ \mathbf b_{0} \\ \mathbf h_l &= \mathbf W_l \left( \mathtt{dropout} \left( \sigma ( \,\mathbf h_{l-1}\, ) \right) \right) + \mathbf b_l\\\end{split}\]

where \(\mathbf x\) is the input tensor, \(\mathbf W_l\) and \(\mathbf b_l\) are the learned weight matrix and bias, respectively, of the \(l\)-th layer, \(\mathbf h_l\) is the hidden representation after layer \(l\), and \(\sigma\) is the activation function.

classmethod build(input_dim, output_dim, hidden_dim=300, n_layers=1, dropout=0.0, activation='relu')

Parameters:

• input_dim (int)

• output_dim (int)

• hidden_dim (int | Sequence[int])

• n_layers (int)

• dropout (float)

• activation (str | torch.nn.Module)

property input_dim: int

Return type:

int

property output_dim: int

Return type:

int

class chemprop.nn.ffn.ConstrainerFFN(n_constraints=1, fp_dim=DEFAULT_HIDDEN_DIM, hidden_dim=300, n_layers=1, dropout=0.0, activation='relu')

Bases: torch.nn.Module, chemprop.nn.hparams.HasHParams, lightning.pytorch.core.mixins.HyperparametersMixin

A ConstrainerFFN adjusts atom or bond property predictions to satisfy molecular constraints by using an MLP to map learned atom or bond embeddings to weights that determine how much of the total adjustment needed is added to each atom or bond prediction.

Parameters:

• n_constraints (int)

• fp_dim (int)

• hidden_dim (int | Sequence[int])

• n_layers (int)

• dropout (float)

• activation (str)

ffn

forward(fp, preds, batch, constraints)

Performs a weighted adjustment to the predictions to satisfy the constraints, with the weights being determined from the learned atom or bond fingerprints via an MLP.

Parameters:

• fp (Tensor) – a tensor of shape b x h containing the atom or bond-level fingerprints, where b is the number of atoms or bonds and h is the length of each fingerprint.

• preds (Tensor) – a tensor of shape b x t containing the atom or bond-level predictions, where t is the number of predictions per atom or bond.

• batch (Tensor) – a tensor of shape b containing indices of which molecule each atom or bond belongs to

• constraints (Tensor) – a tensor of shape m x t containing the values to which the atom or bond-level predictions should sum to for each molecule, where m is the number of molecules in the batch.

Returns:

a tensor of shape b x t containing the atom or bond-level predictions adjusted to satisfy the molecule-level constraints

Return type:

Tensor