Graph Neural Networks for Updating Node/Edge Representations¶
All models based on graph neural networks start with updating node/edge representations. We introduce various GNN models implemented in DGLLifeSci for representation update.
AttentiveFP¶

class
dgllife.model.gnn.attentivefp.
AttentiveFPGNN
(node_feat_size, edge_feat_size, num_layers=2, graph_feat_size=200, dropout=0.0)[source]¶ 
This class performs message passing in AttentiveFP and returns the updated node representations.
 Parameters
node_feat_size (int) – Size for the input node features.
edge_feat_size (int) – Size for the input edge features.
num_layers (int) – Number of GNN layers. Default to 2.
graph_feat_size (int) – Size for the graph representations to be computed. Default to 200.
dropout (float) – The probability for performing dropout. Default to 0.

forward
(g, node_feats, edge_feats)[source]¶ Performs message passing and updates node representations.
 Parameters
g (DGLGraph) – DGLGraph for a batch of graphs.
node_feats (float32 tensor of shape (V, node_feat_size)) – Input node features. V for the number of nodes.
edge_feats (float32 tensor of shape (E, edge_feat_size)) – Input edge features. E for the number of edges.
 Returns
node_feats – Updated node representations.
 Return type
float32 tensor of shape (V, graph_feat_size)
GAT¶

class
dgllife.model.gnn.gat.
GAT
(in_feats, hidden_feats=None, num_heads=None, feat_drops=None, attn_drops=None, alphas=None, residuals=None, agg_modes=None, activations=None)[source]¶ GAT from Graph Attention Networks
 Parameters
in_feats (int) – Number of input node features
hidden_feats (list of int) –
hidden_feats[i]
gives the output size of an attention head in the ith GAT layer.len(hidden_feats)
equals the number of GAT layers. By default, we use[32, 32]
.num_heads (list of int) –
num_heads[i]
gives the number of attention heads in the ith GAT layer.len(num_heads)
equals the number of GAT layers. By default, we use 4 attention heads for each GAT layer.feat_drops (list of float) –
feat_drops[i]
gives the dropout applied to the input features in the ith GAT layer.len(feat_drops)
equals the number of GAT layers. By default, this will be zero for all GAT layers.attn_drops (list of float) –
attn_drops[i]
gives the dropout applied to attention values of edges in the ith GAT layer.len(attn_drops)
equals the number of GAT layers. By default, this will be zero for all GAT layers.alphas (list of float) – Hyperparameters in LeakyReLU, which are the slopes for negative values.
alphas[i]
gives the slope for negative value in the ith GAT layer.len(alphas)
equals the number of GAT layers. By default, this will be 0.2 for all GAT layers.residuals (list of bool) –
residual[i]
decides if residual connection is to be used for the ith GAT layer.len(residual)
equals the number of GAT layers. By default, residual connection is performed for each GAT layer.agg_modes (list of str) – The way to aggregate multihead attention results for each GAT layer, which can be either ‘flatten’ for concatenating allhead results or ‘mean’ for averaging allhead results.
agg_modes[i]
gives the way to aggregate multihead attention results for the ith GAT layer.len(agg_modes)
equals the number of GAT layers. By default, we flatten allhead results for each GAT layer.activations (list of activation function or None) –
activations[i]
gives the activation function applied to the aggregated multihead results for the ith GAT layer.len(activations)
equals the number of GAT layers. By default, no activation is applied for each GAT layer.

forward
(g, feats)[source]¶ Update node representations.
 Parameters
g (DGLGraph) – DGLGraph for a batch of graphs
feats (FloatTensor of shape (N, M1)) –
N is the total number of nodes in the batch of graphs
M1 is the input node feature size, which equals in_feats in initialization
 Returns
feats –
N is the total number of nodes in the batch of graphs
M2 is the output node representation size, which equals hidden_sizes[1] if agg_modes[1] == ‘mean’ and hidden_sizes[1] * num_heads[1] otherwise.
 Return type
FloatTensor of shape (N, M2)
GCN¶

class
dgllife.model.gnn.gcn.
GCN
(in_feats, hidden_feats=None, gnn_norm=None, activation=None, residual=None, batchnorm=None, dropout=None)[source]¶ GCN from SemiSupervised Classification with Graph Convolutional Networks
 Parameters
in_feats (int) – Number of input node features.
hidden_feats (list of int) –
hidden_feats[i]
gives the size of node representations after the ith GCN layer.len(hidden_feats)
equals the number of GCN layers. By default, we use[64, 64]
.gnn_norm (list of str) –
gnn_norm[i]
gives the message passing normalizer for the ith GCN layer, which can be ‘right’, ‘both’ or ‘none’. The ‘right’ normalizer divides the aggregated messages by each node’s indegree. The ‘both’ normalizer corresponds to the symmetric adjacency normalization in the original GCN paper. The ‘none’ normalizer simply sums the messages.len(gnn_norm)
equals the number of GCN layers. By default, we use['none', 'none']
.activation (list of activation functions or None) – If not None,
activation[i]
gives the activation function to be used for the ith GCN layer.len(activation)
equals the number of GCN layers. By default, ReLU is applied for all GCN layers.residual (list of bool) –
residual[i]
decides if residual connection is to be used for the ith GCN layer.len(residual)
equals the number of GCN layers. By default, residual connection is performed for each GCN layer.batchnorm (list of bool) –
batchnorm[i]
decides if batch normalization is to be applied on the output of the ith GCN layer.len(batchnorm)
equals the number of GCN layers. By default, batch normalization is applied for all GCN layers.dropout (list of float) –
dropout[i]
decides the dropout probability on the output of the ith GCN layer.len(dropout)
equals the number of GCN layers. By default, no dropout is performed for all layers.

forward
(g, feats)[source]¶ Update node representations.
 Parameters
g (DGLGraph) – DGLGraph for a batch of graphs
feats (FloatTensor of shape (N, M1)) –
N is the total number of nodes in the batch of graphs
M1 is the input node feature size, which equals in_feats in initialization
 Returns
feats –
N is the total number of nodes in the batch of graphs
M2 is the output node representation size, which equals hidden_sizes[1] in initialization.
 Return type
FloatTensor of shape (N, M2)
MGCN¶

class
dgllife.model.gnn.mgcn.
MGCNGNN
(feats=128, n_layers=3, num_node_types=100, num_edge_types=3000, cutoff=30.0, gap=0.1)[source]¶ MGCN.
MGCN is introduced in Molecular Property Prediction: A Multilevel Quantum Interactions Modeling Perspective.
This class performs message passing in MGCN and returns the updated node representations.
 Parameters
feats (int) – Size for the node and edge embeddings to learn. Default to 128.
n_layers (int) – Number of gnn layers to use. Default to 3.
num_node_types (int) – Number of node types to embed. Default to 100.
num_edge_types (int) – Number of edge types to embed. Default to 3000.
cutoff (float) – Largest center in RBF expansion. Default to 30.
gap (float) – Difference between two adjacent centers in RBF expansion. Default to 0.1.

forward
(g, node_types, edge_dists)[source]¶ Performs message passing and updates node representations.
 Parameters
g (DGLGraph) – DGLGraph for a batch of graphs.
node_types (int64 tensor of shape (V)) – Node types to embed, V for the number of nodes.
edge_dists (float32 tensor of shape (E, 1)) – Distances between end nodes of edges, E for the number of edges.
 Returns
Output node representations.
 Return type
float32 tensor of shape (V, feats * (n_layers + 1))
MPNN¶

class
dgllife.model.gnn.mpnn.
MPNNGNN
(node_in_feats, edge_in_feats, node_out_feats=64, edge_hidden_feats=128, num_step_message_passing=6)[source]¶ MPNN.
MPNN is introduced in Neural Message Passing for Quantum Chemistry.
This class performs message passing in MPNN and returns the updated node representations.
 Parameters
node_in_feats (int) – Size for the input node features.
node_out_feats (int) – Size for the output node representations. Default to 64.
edge_in_feats (int) – Size for the input edge features. Default to 128.
edge_hidden_feats (int) – Size for the hidden edge representations.
num_step_message_passing (int) – Number of message passing steps. Default to 6.

forward
(g, node_feats, edge_feats)[source]¶ Performs message passing and updates node representations.
 Parameters
g (DGLGraph) – DGLGraph for a batch of graphs.
node_feats (float32 tensor of shape (V, node_in_feats)) – Input node features. V for the number of nodes in the batch of graphs.
edge_feats (float32 tensor of shape (E, edge_in_feats)) – Input edge features. E for the number of edges in the batch of graphs.
 Returns
node_feats – Output node representations.
 Return type
float32 tensor of shape (V, node_out_feats)
SchNet¶

class
dgllife.model.gnn.schnet.
SchNetGNN
(node_feats=64, hidden_feats=None, num_node_types=100, cutoff=30.0, gap=0.1)[source]¶ SchNet.
SchNet is introduced in SchNet: A continuousfilter convolutional neural network for modeling quantum interactions.
This class performs message passing in SchNet and returns the updated node representations.
 Parameters
node_feats (int) – Size for node representations to learn. Default to 64.
hidden_feats (list of int) –
hidden_feats[i]
gives the size of hidden representations for the ith interaction layer.len(hidden_feats)
equals the number of interaction layers. Default to[64, 64, 64]
.num_node_types (int) – Number of node types to embed. Default to 100.
cutoff (float) – Largest center in RBF expansion. Default to 30.
gap (float) – Difference between two adjacent centers in RBF expansion. Default to 0.1.

forward
(g, node_types, edge_dists)[source]¶ Performs message passing and updates node representations.
 Parameters
g (DGLGraph) – DGLGraph for a batch of graphs.
node_types (int64 tensor of shape (V)) – Node types to embed, V for the number of nodes.
edge_dists (float32 tensor of shape (E, 1)) – Distances between end nodes of edges, E for the number of edges.
 Returns
node_feats – Updated node representations.
 Return type
float32 tensor of shape (V, node_feats)
Weave¶

class
dgllife.model.gnn.weave.
WeaveGNN
(node_in_feats, edge_in_feats, num_layers=2, hidden_feats=50, activation=<function relu>)[source]¶ The component of Weave for updating node and edge representations.
Weave is introduced in Molecular Graph Convolutions: Moving Beyond Fingerprints.
 Parameters
node_in_feats (int) – Size for the input node features.
edge_in_feats (int) – Size for the input edge features.
num_layers (int) – Number of Weave layers to use, which is equivalent to the times of message passing. Default to 2.
hidden_feats (int) – Size for the hidden node and edge representations. Default to 50.
activation (callable) – Activation function to be used. It cannot be None. Default to ReLU.

forward
(g, node_feats, edge_feats, node_only=True)[source]¶ Updates node representations (and edge representations).
 Parameters
g (DGLGraph) – DGLGraph for a batch of graphs.
node_feats (float32 tensor of shape (V, node_in_feats)) – Input node features. V for the number of nodes in the batch of graphs.
edge_feats (float32 tensor of shape (E, edge_in_feats)) – Input edge features. E for the number of edges in the batch of graphs.
node_only (bool) – Whether to return updated node representations only or to return both node and edge representations. Default to True.
 Returns
float32 tensor of shape (V, gnn_hidden_feats) – Updated node representations.
float32 tensor of shape (E, gnn_hidden_feats), optional – This is returned only when
node_only==False
. Updated edge representations.
GIN¶

class
dgllife.model.gnn.gin.
GIN
(num_node_emb_list, num_edge_emb_list, num_layers=5, emb_dim=300, JK='last', dropout=0.5)[source]¶ Graph Isomorphism Network from Strategies for Pretraining Graph Neural Networks
This module is for updating node representations only.
 Parameters
num_node_emb_list (list of int) – num_node_emb_list[i] gives the number of items to embed for the ith categorical node feature variables. E.g. num_node_emb_list[0] can be the number of atom types and num_node_emb_list[1] can be the number of atom chirality types.
num_edge_emb_list (list of int) – num_edge_emb_list[i] gives the number of items to embed for the ith categorical edge feature variables. E.g. num_edge_emb_list[0] can be the number of bond types and num_edge_emb_list[1] can be the number of bond direction types.
num_layers (int) – Number of GIN layers to use. Default to 5.
emb_dim (int) – The size of each embedding vector. Default to 300.
JK (str) –
JK for jumping knowledge as in Representation Learning on Graphs with Jumping Knowledge Networks. It decides how we are going to combine the alllayer node representations for the final output. There can be four options for this argument,
concat
,last
,max
andsum
. Default to ‘last’.'concat'
: concatenate the output node representations from all GIN layers'last'
: use the node representations from the last GIN layer'max'
: apply max pooling to the node representations across all GIN layers'sum'
: sum the output node representations from all GIN layers
dropout (float) – Dropout to apply to the output of each GIN layer. Default to 0.5

forward
(g, categorical_node_feats, categorical_edge_feats)[source]¶ Update node representations
 Parameters
g (DGLGraph) – DGLGraph for a batch of graphs
categorical_node_feats (list of LongTensor of shape (N)) –
Input categorical node features
len(categorical_node_feats) should be the same as len(self.node_embeddings)
N is the total number of nodes in the batch of graphs
categorical_edge_feats (list of LongTensor of shape (E)) –
Input categorical edge features
len(categorical_edge_feats) should be the same as len(num_edge_emb_list) in the arguments
E is the total number of edges in the batch of graphs
 Returns
final_node_feats – Output node representations, N for the number of nodes and M for output size. In particular, M will be emb_dim * (num_layers + 1) if self.JK == ‘concat’ and emb_dim otherwise.
 Return type
float32 tensor of shape (N, M)
GraphSAGE¶

class
dgllife.model.gnn.graphsage.
GraphSAGE
(in_feats, hidden_feats=None, activation=None, dropout=None, aggregator_type=None)[source]¶ GraphSAGE from Inductive Representation Learning on Large Graphs
 Parameters
in_feats (int) – Number of input node features.
hidden_feats (list of int) –
hidden_feats[i]
gives the size of node representations after the ith GraphSAGE layer.len(hidden_feats)
equals the number of GraphSAGE layers. By default, we use[64, 64]
.activation (list of activation functions or None) – If not None,
activation[i]
gives the activation function to be used for the ith GraphSAGE layer.len(activation)
equals the number of GraphSAGE layers. By default, ReLU is applied for all GraphSAGE layers.dropout (list of float or None) –
dropout[i]
decides the dropout probability on the output of the ith GraphSAGE layer.len(dropout)
equals the number of GraphSAGE layers. By default, no dropout is performed for all layers.aggregator_type (list of str) –
aggregator_type[i]
decides the aggregator type for the ith GraphSAGE layer, which can be one of'mean'
,'gcn'
,'pool'
,'lstm'
. By default, we use'mean'
for all layers.

forward
(g, feats)[source]¶ Update node representations.
 Parameters
g (DGLGraph) – DGLGraph for a batch of graphs
feats (FloatTensor of shape (N, M1)) –
N is the total number of nodes in the batch of graphs
M1 is the input node feature size, which equals in_feats in initialization
 Returns
feats –
N is the total number of nodes in the batch of graphs
M2 is the output node representation size, which equals hidden_sizes[1] in initialization.
 Return type
FloatTensor of shape (N, M2)
WLN¶

class
dgllife.model.gnn.wln.
WLN
(node_in_feats, edge_in_feats, node_out_feats=300, n_layers=3, project_in_feats=True, set_comparison=True)[source]¶ WeisfeilerLehman Network (WLN)
WLN is introduced in Predicting Organic Reaction Outcomes with WeisfeilerLehman Network.
This class performs message passing and updates node representations.
 Parameters
node_in_feats (int) – Size for the input node features.
edge_in_feats (int) – Size for the input edge features.
node_out_feats (int) – Size for the output node representations. Default to 300.
n_layers (int) – Number of times for message passing. Note that same parameters are shared across n_layers message passing. Default to 3.
project_in_feats (bool) – Whether to project input node features. If this is False, we expect node_in_feats to be the same as node_out_feats. Default to True.
set_comparison (bool) – Whether to perform final node representation update mimicking set comparison. Default to True.

forward
(g, node_feats, edge_feats)[source]¶ Performs message passing and updates node representations.
 Parameters
g (DGLGraph) – DGLGraph for a batch of graphs
node_feats (float32 tensor of shape (V, node_in_feats)) – Input node features. V for the number of nodes.
edge_feats (float32 tensor of shape (E, edge_in_feats)) – Input edge features. E for the number of edges.
 Returns
Updated node representations.
 Return type
float32 tensor of shape (V, node_out_feats)
GNNOGB¶

class
dgllife.model.gnn.gnn_ogb.
GNNOGB
(in_edge_feats, num_node_types=1, hidden_feats=300, n_layers=5, batchnorm=True, activation=<function relu>, dropout=0.0, gnn_type='gcn', virtual_node=True, residual=False, jk=False)[source]¶ Variant of GCN/GIN from Open Graph Benchmark: Datasets for Machine Learning on Graphs
 Parameters
in_edge_feats (int) – Number of input edge features.
num_node_types (int) – Number of node types to embed. (Default: 1)
hidden_feats (int) – Size for hidden representations. (Default: 300)
n_layers (int) – Number of GNN layers to use. (Default: 5)
batchnorm (bool) – Whether to apply batch normalization. (Default: True)
activation (callable or None) – Activation function to apply to the output of each GNN layer except for the last layer. If None, no activation will be applied. (Default: ReLU)
dropout (float) – The probability for dropout. (Default: 0, i.e. no dropout)
gnn_type (str) – The GNN type to use, which can be either ‘gcn’ or ‘gin’. (Default: ‘gcn’)
virtual_node (bool) – Whether to use virtual node. (Default: True)
residual (bool) – Whether to apply residual connections for virtual node embeddings. (Default: False)
jk (bool) – Whether to sum over the output of all GNN layers as in JK networks. (Default: False)

forward
(g, node_feats, edge_feats)[source]¶ Update node representations.
 Parameters
g (DGLGraph) – DGLGraph for a batch of graphs
node_feats (LongTensor of shape (N, 1)) – Input categorical node features. N for the number of nodes.
edge_feats (FloatTensor of shape (E, in_edge_feats)) – Input edge features. E for the number of edges.
 Returns
Output node representations
 Return type
FloatTensor of shape (N, hidden_feats)