CuGraphGATConv

class dgl.nn.pytorch.conv.CuGraphGATConv(in_feats, out_feats, num_heads, feat_drop=0.0, negative_slope=0.2, residual=False, activation=None, bias=True)

Bases: CuGraphBaseConv

Graph attention layer from Graph Attention Networks, with the sparse aggregation accelerated by cugraph-ops.

See dgl.nn.pytorch.conv.GATConv for mathematical model.

This module depends on pylibcugraphops package, which can be installed via conda install -c nvidia pylibcugraphops=23.04. pylibcugraphops 23.04 requires python 3.8.x or 3.10.x.

Note

This is an experimental feature.

Parameters:

• in_feats (int) – Input feature size.

• out_feats (int) – Output feature size.

• num_heads (int) – Number of heads in Multi-Head Attention.

• feat_drop (float, optional) – Dropout rate on feature. Defaults: 0.

• negative_slope (float, optional) – LeakyReLU angle of negative slope. Defaults: 0.2.

• residual (bool, optional) – If True, use residual connection. Defaults: False.

• activation (callable activation function/layer or None, optional.) – If not None, applies an activation function to the updated node features. Default: None.

• bias (bool, optional) – If True, learns a bias term. Defaults: True.

Examples

>>> import dgl
>>> import torch
>>> from dgl.nn import CuGraphGATConv
>>> device = 'cuda'
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3])).to(device)
>>> g = dgl.add_self_loop(g)
>>> feat = torch.ones(6, 10).to(device)
>>> conv = CuGraphGATConv(10, 2, num_heads=3).to(device)
>>> res = conv(g, feat)
>>> res
tensor([[[ 0.2340,  1.9226],
        [ 1.6477, -1.9986],
        [ 1.1138, -1.9302]],
        [[ 0.2340,  1.9226],
        [ 1.6477, -1.9986],
        [ 1.1138, -1.9302]],
        [[ 0.2340,  1.9226],
        [ 1.6477, -1.9986],
        [ 1.1138, -1.9302]],
        [[ 0.2340,  1.9226],
        [ 1.6477, -1.9986],
        [ 1.1138, -1.9302]],
        [[ 0.2340,  1.9226],
        [ 1.6477, -1.9986],
        [ 1.1138, -1.9302]],
        [[ 0.2340,  1.9226],
        [ 1.6477, -1.9986],
        [ 1.1138, -1.9302]]], device='cuda:0', grad_fn=<ViewBackward0>)

forward(g, feat, max_in_degree=None)

Forward computation.

Parameters:

• g (DGLGraph) – The graph.

• feat (torch.Tensor) – Input features of shape \((N, D_{in})\).

• max_in_degree (int) – Maximum in-degree of destination nodes. It is only effective when g is a DGLBlock, i.e., bipartite graph. When g is generated from a neighbor sampler, the value should be set to the corresponding fanout. If not given, max_in_degree will be calculated on-the-fly.

Returns:

The output feature of shape \((N, H, D_{out})\) where \(H\) is the number of heads, and \(D_{out}\) is size of output feature.

Return type:

torch.Tensor

reset_parameters()

Reinitialize learnable parameters.