GATConv

class graphstorm.model.GATConv(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

Graph attention layer from Graph Attention Network.

The message passing formulas of GATConv are:

\[h_i^{(l+1)} = \sum_{j\in \mathcal{N}(i)} \alpha_{i,j} W^{(l)} 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} &= \mathrm{LeakyReLU}\left(\vec{a}^T [W h_{i} \| W h_{j}]\right)\end{aligned}\end{align} \]

Note:

• GATEConv is only effective on homogeneous graphs.

Examples:

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

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

Parameters

in_feat: int

Input feature size.

out_feat: int

Output feature size.

num_heads: int

Number of heads in Multi-head attention.

activation: callable, optional

Activation function. Default: relu

dropout: float, optional

Dropout rate. Default: 0.0

bias: bool, 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)

GAT 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.