GATv2Conv

class graphstorm.model.GATv2Conv(in_feat, out_feat, num_heads, activation=<function relu>, dropout=0.0, bias=True, num_ffn_layers_in_gnn=0, ffn_activation=<function relu>)

Bases: Module

GATv2 Convolutional layer from How Attentive are Graph Attention Networks?.

The message passing formulas of GATv2Conv are:

\[h_i^{(l+1)} = \sum_{j\in \mathcal{N}(i)} \alpha_{ij}^{(l)} W^{(l)}_{right} h_j^{(l)}\]

where \(\alpha_{ij}\) is the attention score bewteen node \(i\) and node \(j\):

\[ \begin{align}\begin{aligned}\alpha_{ij}^{(l)} &= \mathrm{softmax_i} (e_{ij}^{(l)})\\e_{ij}^{(l)} &= {\vec{a}^T}^{(l)}\mathrm{LeakyReLU}\left( W^{(l)}_{left} h_{i} + W^{(l)}_{right} h_{j}\right)\end{aligned}\end{align} \]

Note:

  • GATv2Conv is only effective on homogeneous graphs.

Examples:

# suppose graph and input_feature are ready
from graphstorm.model import GATv2Conv

layer = GATv2Conv(h_dim, h_dim, num_heads, num_ffn_layers_in_gnn)
h = layer(g, input_feature)

Parameters

in_featint

Input feature size.

out_featint

Output feature size.

num_headsint

Number of heads in Multi-head attention.

activationcallable, optional

Activation function. Default: relu

dropoutfloat, optional

Dropout rate. Default: 0.0

biasbool, optional

True if bias is added. Default: True

num_ffn_layers_in_gnn: int, optional

Number of layers of ngnn between gnn layers. Default: 0

ffn_actication: torch.nn.functional, optional

Activation Method for ngnn. Default: relu

forward(g, inputs)

GATv2 layer Forward computation.

Parameters

g: DGLHeteroGraph

Input DGL heterogenous graph.

inputs: dict of Tensor

Node features for the default node type in the format of {dgl.DEFAULT_NTYPE: tensor}. The definition of dgl.DEFAULT_NTYPE can be found at DGL official Github site.

Returns

dict of Tensor: New node embeddings for the default node type in the format of

{dgl.DEFAULT_NTYPE: tensor}. The definition of dgl.DEFAULT_NTYPE can be found at DGL official Github site.