3. pahelix.networks

3.1. basic_block

Some frequently used basic blocks

class pahelix.networks.basic_block.Activation(*args: Any, **kwargs: Any)

forward(x)

tbd

class pahelix.networks.basic_block.MLP(*args: Any, **kwargs: Any)

forward(x)

Parameters:

x (tensor) – (-1, dim).

3.2. compound_encoder

Basic Encoder for compound atom/bond features.

class pahelix.networks.compound_encoder.AtomEmbedding(*args: Any, **kwargs: Any)

Atom Encoder

forward(node_features)

Parameters:

node_features (dict of tensor) – node features.

class pahelix.networks.compound_encoder.AtomFloatEmbedding(*args: Any, **kwargs: Any)

Atom Float Encoder

forward(feats)

Parameters:

feats (dict of tensor) – node float features.

class pahelix.networks.compound_encoder.BondAngleFloatRBF(*args: Any, **kwargs: Any)

Bond Angle Float Encoder using Radial Basis Functions

forward(bond_angle_float_features)

Parameters:

bond_angle_float_features (dict of tensor) – bond angle float features.

class pahelix.networks.compound_encoder.BondEmbedding(*args: Any, **kwargs: Any)

Bond Encoder

forward(edge_features)

Parameters:

edge_features (dict of tensor) – edge features.

class pahelix.networks.compound_encoder.BondFloatRBF(*args: Any, **kwargs: Any)

Bond Float Encoder using Radial Basis Functions

forward(bond_float_features)

Parameters:

bond_float_features (dict of tensor) – bond float features.

3.3. gnn_block

Blocks for Graph Neural Network (GNN)

Adapted from https://github.com/snap-stanford/pretrain-gnns/blob/master/chem/model.py

class pahelix.networks.gnn_block.GIN(*args: Any, **kwargs: Any)

Implementation of Graph Isomorphism Network (GIN) layer with edge features

forward(graph, node_feat, edge_feat)

Parameters:

• node_feat (tensor) – node features with shape (num_nodes, feature_size).

• edge_feat (tensor) – edges features with shape (num_edges, feature_size).

class pahelix.networks.gnn_block.GraphNorm(*args: Any, **kwargs: Any)

Implementation of graph normalization. Each node features is divied by sqrt(num_nodes) per graphs.

Parameters:

• graph – the graph object from (Graph)

• feature – A tensor with shape (num_nodes, feature_size).

Returns:

A tensor with shape (num_nodes, hidden_size)

References:

[1] BENCHMARKING GRAPH NEURAL NETWORKS. https://arxiv.org/abs/2003.00982

forward(graph, feature)

graph norm

class pahelix.networks.gnn_block.MeanPool(*args: Any, **kwargs: Any)

TODO: temporary class due to pgl mean pooling

forward(graph, node_feat)

mean pooling

3.4. involution_block

class pahelix.networks.involution_block.Involution2D(in_channel, out_channel, sigma_mapping=None, kernel_size=7, stride=1, groups=1, reduce_ratio=1, dilation=1, padding=3)

Involution module.

Parameters:

• in_channel – The channel size of input.

• out_channel – The channel size of output.

• sigma_mapping – Sigma mapping.

• kernel_size – Kernel size.

• stride – Stride size.

• groups – Group size.

• reduce_ratio – The ratio of reduce.

• dilation – The dilation size.

• padding – The padding size.

Returns:

Tbe output of Involution2D block.

Return type:

output

References:

[1] Involution: Inverting the Inherence of Convolution for Visual Recognition. https://arxiv.org/abs/2103.06255

forward(x)

Involution block

3.5. lstm_block

Lstm block.

pahelix.networks.lstm_block.lstm_encoder(input, hidden_size, n_layer=1, is_bidirectory=True, param_initializer=None, name='lstm')

The encoder is composed of a stack of lstm layers.

Parameters:

• input – The input of lstm encoder.

• hidden_size – The hidden size of lstm.

• n_layer – The number of lstm layers.

• is_bidirectory – True if the lstm is bidirectory.

• param_initializer – The parameter initializer for lstm encoder.

• name – The prefix of the parameters’ name in lstm encoder.

Returns:

The hidden units of lstm encoder. checkpoints: The checkpoints for recompute mechanism.

Return type:

hidden

3.6. optimizer

class pahelix.networks.optimizer.AdamW(*args, **kwargs)

AdamW object for dygraph.

apply_optimize(loss, startup_program, params_grads)

Update params with weight decay.

3.7. pre_post_process

pahelix.networks.pre_post_process.pre_post_process_layer(prev_out, out, process_cmd, dropout_rate=0.0, epsilon=1e-05, name='', is_test=False)

Add residual connection, layer normalization and droput to the out tensor optionally according to the value of process_cmd.

This will be used before or after multi-head attention and position-wise feed-forward networks.

3.8. resnet_block

Resnet block.

pahelix.networks.resnet_block.resnet_encoder(input, hidden_size, n_layer=1, filter_size=3, act='gelu', epsilon=1e-06, param_initializer=None, name='resnet')

The encoder is composed of a stack of resnet layers.

Parameters:

• input – The input of resnet encoder.

• hidden_size – The hidden size of resnet.

• n_layer – The number of resnet layers.

• act – The activation function.

• param_initializer – The parameter initializer for resnet encoder.

• name – The prefix of the parameters’ name in resnet encoder.

Returns:

The hidden units of resnet encoder. checkpoints: The checkpoints for recompute mechanism.

Return type:

hidden

3.9. transformer_block

Transformer block.

pahelix.networks.transformer_block.multi_head_attention(queries, keys, values, attn_bias, d_key, d_value, d_model, n_head=1, dropout_rate=0.0, cache=None, gather_idx=None, store=False, param_initializer=None, lr=1.0, name='multi_head_att', is_test=False)

Multi-Head Attention.

Note that attn_bias is added to the logit before computing softmax activiation to mask certain selected positions so that they will not considered in attention weights.

pahelix.networks.transformer_block.positionwise_feed_forward(x, d_inner_hid, d_hid, dropout_rate, hidden_act, num_flatten_dims=2, param_initializer=None, name='ffn', is_test=False)

Position-wise Feed-Forward Networks.

This module consists of two linear transformations with a ReLU activation in between, which is applied to each position separately and identically.

pahelix.networks.transformer_block.transformer_encoder(enc_input, attn_bias, n_layer, n_head, d_key, d_value, d_model, d_inner_hid, prepostprocess_dropout, attention_dropout, act_dropout, hidden_act, preprocess_cmd='n', postprocess_cmd='da', param_initializer=None, name='', epsilon=1e-05, n_layer_per_block=1, param_share='normal', caches=None, gather_idx=None, store=False, is_test=False)

The encoder is composed of a stack of identical layers returned by calling transformer_encoder_layer.

pahelix.networks.transformer_block.transformer_encoder_layer(input, attn_bias, n_head, d_key, d_value, d_model, d_inner_hid, prepostprocess_dropout, attention_dropout, act_dropout, hidden_act, preprocess_cmd='n', postprocess_cmd='da', param_initializer=None, name='', epsilon=1e-05, cache=None, gather_idx=None, store=False, is_test=False)

The encoder layers that can be stacked to form a deep encoder.

This module consits of a multi-head (self) attention followed by position-wise feed-forward networks and both the two components companied with the pre_process_layer / post_process_layer to add residual connection, layer normalization and droput.

3.10. Helpful Link

Please refer to our GitHub repo to see the whole module.