Designing a GCN
Designing a GCN with NAGL primarily involves creating an instance of the [ModelConfig
] class.
In Python
A ModelConfig class can be created can be done straightforwardly in Python.
from openff.nagl.features import atoms, bonds
from openff.nagl import GNNModel
from openff.nagl.nn import gcn
from openff.nagl.nn.activation import ActivationFunction
from openff.nagl.nn import postprocess
from openff.nagl.config.model import (
ModelConfig,
ConvolutionModule, ReadoutModule,
ConvolutionLayer, ForwardLayer,
)
First we can specify our desired features. These should be instances of the feature classes.
atom_features = (
atoms.AtomicElement(["C", "H", "O", "N", "P", "S"]),
atoms.AtomConnectivity(),
...
)
bond_features = (
bonds.BondOrder(),
...
)
Next, we can design convolution and readout modules. For example:
convolution_module = ConvolutionModule(
architecture="SAGEConv",
# construct 2 layers with dropout 0 (default),
# hidden feature size 512, and ReLU activation function
# these layers can also be individually specified,
# but we just duplicate the layer 6 times for identical layers
layers=[
ConvolutionLayer(
hidden_feature_size=512,
activation_function="ReLU",
aggregator_type="mean"
)
] * 2,
)
# define our readout module/s
# multiple are allowed but let's focus on charges
readout_modules = {
# key is the name of output property, any naming is allowed
"charges": ReadoutModule(
pooling="atoms",
postprocess="compute_partial_charges",
# 1 layers
layers=[
ForwardLayer(
hidden_feature_size=512,
activation_function="ReLU",
)
],
)
}
We can now bring it all together as a ModelConfig
and create a GNNModel
.
model_config = ModelConfig(
version="0.1",
atom_features=atom_features,
bond_features=bond_features,
convolution=convolution_module,
readouts=readout_modules,
)
model = GNNModel(model_config)
From YAML
Or if you prefer, the same model architecture can be specified as a YAML file:
version: '0.1'
convolution:
architecture: SAGEConv
layers:
- hidden_feature_size: 512
activation_function: ReLU
dropout: 0
aggregator_type: mean
- hidden_feature_size: 512
activation_function: ReLU
dropout: 0
aggregator_type: mean
readouts:
charges:
pooling: atoms
postprocess: compute_partial_charges
layers:
- hidden_feature_size: 128
activation_function: Sigmoid
dropout: 0
atom_features:
- name: atomic_element
categories: ["C", "H", "O", "N", "P", "S"]
- name: atom_connectivity
categories: [1, 2, 3, 4, 5, 6]
- name: atom_hybridization
- name: atom_in_ring_of_size
ring_size: 3
- name: atom_in_ring_of_size
ring_size: 4
- name: atom_in_ring_of_size
ring_size: 5
- name: atom_in_ring_of_size
ring_size: 6
bond_features:
- name: bond_is_in_ring
And then loaded into a config using the [ModelConfig.from_yaml()
] method:
from openff.nagl import GNNModel
from openff.nagl.config import ModelConfig
model = GNNModel(ModelConfig.from_yaml("model.yaml"))
Here we’ll go through each option, what it means, and where to find the available choices.
atom_features
and bond_features
These arguments specify the featurization scheme for the model (see Featurization). atom_features
takes a tuple of features from the openff.nagl.features.atoms
module, and bond_features
a tuple of features from the openff.nagl.features.bonds
module. Each feature is a class that must be instantiated, possibly with some arguments. Custom features may be implemented by subclassing AtomFeature
or BondFeature
; both share the interface of their base class Feature
.
convolution_architecture
The convolution_architecture
argument specifies the structure of the convolution module. Available options are provided in the [openff.nagl.config.model
] module.
Number of Features and Layers
Each module comprises a number of layers that must be individually specified.
For example, a [ConvolutionModule
] consists of specified [ConvolutionLayer
]s. A [ReadoutModule
] consists of specified [ForwardLayer
]s.
The “convolution” arguments define the update network in the convolution module, and the “readout” the network in the readout module (see The Convolution Module: Message-passing Graph Convolutional Networks and The Readout Module: The Charge Equilibration Method). Read the GNNModel
docstring carefully to determine which layers are considered hidden.
activation_function
The activation_function
argument defines the activation function used by the readout network (see Neural Networks - a quick primer). The activation function used by the convolution network is currently fixed to ReLU. Available activation functions are the variants (attributes) of the ActivationFunction
class. When using the Python API rather than the YAML interface, other activation functions from PyTorch can be used.
postprocess_layer
postprocess_layer
specifies a PyTorch Module
that performs post-processing in the readout module (see The Readout Module: The Charge Equilibration Method). Post-processing layers inherit from the PostprocessLayer
class and are found in the openff.nagl.nn.postprocess
module.