CuGraphRelGraphConv

class dgl.nn.pytorch.conv.CuGraphRelGraphConv(in_feat, out_feat, num_rels, regularizer=None, num_bases=None, bias=True, activation=None, self_loop=True, dropout=0.0, layer_norm=False, max_in_degree=None)[source]

Bases: torch.nn.modules.module.Module

An accelerated relational graph convolution layer from Modeling Relational Data with Graph Convolutional Networks that leverages the highly-optimized aggregation primitives in cugraph-ops.

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

This module depends on pylibcugraphops package, which can be installed via conda install -c nvidia pylibcugraphops>=22.12.

Note

This is an experimental feature. Compared with dgl.nn.pytorch.conv.RelGraphConv, this model:

  • Only works on cuda devices.

  • Only supports basis-decomposition regularization.

Parameters

  • in_feat (int) – Input feature size.

  • out_feat (int) – Output feature size.

  • num_rels (int) – Number of relations.

  • regularizer (str, optional) –

    Which weight regularizer to use (“basis” or None):

    • ”basis” is for basis-decomposition.

    • None applies no regularization.

    Default: None.

  • num_bases (int, optional) – Number of bases. It comes into effect when a regularizer is applied. Default: None.

  • bias (bool, optional) – True if bias is added. Default: True.

  • activation (callable, optional) – Activation function. Default: None.

  • self_loop (bool, optional) – True to include self loop message. Default: True.

  • dropout (float, optional) – Dropout rate. Default: 0.0.

  • layer_norm (bool, optional) – True to add layer norm. Default: False.

  • max_in_degree (int, optional) – Maximum number of sampled neighbors of a destination node, i.e. maximum in degree of destination nodes. If None, it will be calculated on the fly during forward().

Examples

>>> import dgl
>>> import torch as th
>>> from dgl.nn import CuGraphRelGraphConv
...
>>> device = 'cuda'
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3])).to(device)
>>> feat = th.ones(6, 10).to(device)
>>> conv = CuGraphRelGraphConv(
...     10, 2, 3, regularizer='basis', num_bases=2).to(device)
>>> etype = th.tensor([0,1,2,0,1,2]).to(device)
>>> res = conv(g, feat, etype)
>>> res
tensor([[-1.7774, -2.0184],
        [-1.4335, -2.3758],
        [-1.7774, -2.0184],
        [-0.4698, -3.0876],
        [-1.4335, -2.3758],
        [-1.4331, -2.3295]], device='cuda:0', grad_fn=<AddBackward0>)
forward(g, feat, etypes, norm=None)[source]

Forward computation.

Parameters

  • g (DGLGraph) – The graph.

  • feat (torch.Tensor) – A 2D tensor of node features. Shape: \((|V|, D_{in})\).

  • etypes (torch.Tensor) – A 1D integer tensor of edge types. Shape: \((|E|,)\). Note that cugraph-ops only accepts edge type tensors in int32, so any input of other integer types will be casted into int32, thus introducing some overhead. Pass in int32 tensors directly for best performance.

  • norm (torch.Tensor, optional) – A 1D tensor of edge norm value. Shape: \((|E|,)\).

Returns

New node features. Shape: \((|V|, D_{out})\).

Return type

torch.Tensor