- class openff.nagl.GNNModel(convolution_architecture: Union[str, BaseGCNStack], n_convolution_hidden_features: int, n_convolution_layers: int, n_readout_hidden_features: int, n_readout_layers: int, activation_function: Union[str, ActivationFunction], postprocess_layer: Union[str, PostprocessLayer], readout_name: str, learning_rate: float, atom_features: Tuple[AtomFeature, ...], bond_features: Tuple[BondFeature, ...], loss_function: Callable = rmse_loss, convolution_dropout: float = 0, readout_dropout: float = 0)[source]
Bases:
BaseGNNModel
A model that applies a graph convolutional step followed by pooling and readout steps.
- Parameters
convolution_architecture (
Union[str
,BaseGCNStack]
) – The graph convolution architecture. This can be given either as a class, e.g.SAGEConvStack
or as a string, e.g."SAGEConv"
.n_convolution_hidden_features (
int
) – The number of features in each of the hidden convolutional layers.n_convolution_layers (
int
) – The number of hidden convolutional layers to generate. These are the layers in the convolutional module between the input layer and the pooling layer.n_readout_hidden_features (
int
) – The number of features in each of the hidden readout layers.n_readout_layers (
int
) – The number of hidden readout layers to generate. These are the layers between the convolution module’s pooling layer and the readout module’s output layer. The pooling layer may be considered to be both the convolution module’s output layer and the readout module’s input layer.activation_function (
Union[str
,ActivationFunction]
) – The activation function to use for the readout module. This can be given either as a class, e.g.ReLU
, or as a string, e.g."ReLU"
.postprocess_layer (
Union[str
,PostprocessLayer]
) – The postprocess layer to use. This can be given either as a class, e.g.ComputePartialCharges
, or as a string, e.g."compute_partial_charges"
.readout_name (
str
) – A human-readable name for the readout module.learning_rate (
float
) – The learning rate for optimization.atom_features (
Tuple[AtomFeature
,]
) – The atom features to use.bond_features (
Tuple[BondFeature
,]
) – The bond features to use.loss_function (
Callable[[torch.Tensor
,torch.Tensor]
,torch.Tensor]
) – The loss function. This is RMSE by default, but can be any function that takes a predicted and target tensor and returns a scalar loss tensor.convolution_dropout (
float
) – The dropout probability to use in the convolutional layers.readout_dropout (
float
) – The dropout probability to use in the readout layers.
Methods
Compute the trained property for a molecule.
Construct a
GNNModel
from a YAML fileLoad a model from a file.
Save this model to a file.
Attributes
The number of features used to represent an atom
- compute_property(molecule: Molecule, as_numpy: bool = False) Tensor [source]
Compute the trained property for a molecule.
- Parameters
- Returns
result (
torch.Tensor
ornumpy.ndarray
)
- classmethod from_yaml_file(*paths, **kwargs) GNNModel [source]
Construct a
GNNModel
from a YAML file
- classmethod load(model: str, eval_mode: bool = True)[source]
Load a model from a file.
- Parameters
- Returns
model (
GNNModel
)
Examples
>>> model.save("model.pt") >>> new_model = GNNModel.load("model.pt")
Notes
This method is not compatible with normal Pytorch models saved with
torch.save
, as it expects a dictionary of hyperparameters and a state dictionary.
- property n_atom_features: int
The number of features used to represent an atom
- save(path: str)[source]
Save this model to a file.
- Parameters
path (
str
) – The path to save this file to.
Examples
>>> model.save("model.pt") >>> new_model = GNNModel.load("model.pt")
Notes
This method writes a dictionary of the hyperparameters and the state dictionary, with the keys “hyperparameters” and “state_dict”.