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)[source]ΒΆ Bases:
dgl.nn.pytorch.conv.cugraph_base.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 viaconda 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)[source]ΒΆ 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 aDGLBlock
, i.e., bipartite graph. Wheng
is generated from a neighbor sampler, the value should be set to the correspondingfanout
. 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