layers.gcc_module

Module Contents

Classes

GATLayer

SELayer

Squeeze-and-excitation networks

ApplyNodeFunc

Update the node feature hv with MLP, BN and ReLU.

MLP

MLP with linear output

UnsupervisedGAT

UnsupervisedMPNN

MPNN from

UnsupervisedGIN

GIN model

GraphEncoder

MPNN from

class layers.gcc_module.GATLayer(g, in_dim, out_dim)[source]

Bases: torch.nn.Module

edge_attention(self, edges)[source]
message_func(self, edges)[source]
reduce_func(self, nodes)[source]
forward(self, h)[source]
class layers.gcc_module.SELayer(in_channels, se_channels)[source]

Bases: torch.nn.Module

Squeeze-and-excitation networks

forward(self, x)[source]
class layers.gcc_module.ApplyNodeFunc(mlp, use_selayer)[source]

Bases: torch.nn.Module

Update the node feature hv with MLP, BN and ReLU.

forward(self, h)[source]
class layers.gcc_module.MLP(num_layers, input_dim, hidden_dim, output_dim, use_selayer)[source]

Bases: torch.nn.Module

MLP with linear output

forward(self, x)[source]
class layers.gcc_module.UnsupervisedGAT(node_input_dim, node_hidden_dim, edge_input_dim, num_layers, num_heads)[source]

Bases: torch.nn.Module

forward(self, g, n_feat, e_feat)[source]
class layers.gcc_module.UnsupervisedMPNN(output_dim=32, node_input_dim=32, node_hidden_dim=32, edge_input_dim=32, edge_hidden_dim=32, num_step_message_passing=6, lstm_as_gate=False)[source]

Bases: torch.nn.Module

MPNN from Neural Message Passing for Quantum Chemistry

node_input_dimint

Dimension of input node feature, default to be 15.

edge_input_dimint

Dimension of input edge feature, default to be 15.

output_dimint

Dimension of prediction, default to be 12.

node_hidden_dimint

Dimension of node feature in hidden layers, default to be 64.

edge_hidden_dimint

Dimension of edge feature in hidden layers, default to be 128.

num_step_message_passingint

Number of message passing steps, default to be 6.

num_step_set2setint

Number of set2set steps

num_layer_set2setint

Number of set2set layers

forward(self, g, n_feat, e_feat)[source]

Predict molecule labels

gDGLGraph

Input DGLGraph for molecule(s)

n_feattensor of dtype float32 and shape (B1, D1)

Node features. B1 for number of nodes and D1 for the node feature size.

e_feattensor of dtype float32 and shape (B2, D2)

Edge features. B2 for number of edges and D2 for the edge feature size.

res : Predicted labels

class layers.gcc_module.UnsupervisedGIN(num_layers, num_mlp_layers, input_dim, hidden_dim, output_dim, final_dropout, learn_eps, graph_pooling_type, neighbor_pooling_type, use_selayer)[source]

Bases: torch.nn.Module

GIN model

forward(self, g, h, efeat)[source]
class layers.gcc_module.GraphEncoder(positional_embedding_size=32, max_node_freq=8, max_edge_freq=8, max_degree=128, freq_embedding_size=32, degree_embedding_size=32, output_dim=32, node_hidden_dim=32, edge_hidden_dim=32, num_layers=6, num_heads=4, num_step_set2set=6, num_layer_set2set=3, norm=False, gnn_model='mpnn', degree_input=False, lstm_as_gate=False)[source]

Bases: torch.nn.Module

MPNN from Neural Message Passing for Quantum Chemistry

node_input_dimint

Dimension of input node feature, default to be 15.

edge_input_dimint

Dimension of input edge feature, default to be 15.

output_dimint

Dimension of prediction, default to be 12.

node_hidden_dimint

Dimension of node feature in hidden layers, default to be 64.

edge_hidden_dimint

Dimension of edge feature in hidden layers, default to be 128.

num_step_message_passingint

Number of message passing steps, default to be 6.

num_step_set2setint

Number of set2set steps

num_layer_set2setint

Number of set2set layers

forward(self, g, return_all_outputs=False)[source]

Predict molecule labels

gDGLGraph

Input DGLGraph for molecule(s)

n_feattensor of dtype float32 and shape (B1, D1)

Node features. B1 for number of nodes and D1 for the node feature size.

e_feattensor of dtype float32 and shape (B2, D2)

Edge features. B2 for number of edges and D2 for the edge feature size.

res : Predicted labels