CuGraphSAGEConv

class dgl.nn.pytorch.conv.CuGraphSAGEConv(in_feats, out_feats, aggregator_type='mean', feat_drop=0.0, bias=True)[source]

基础类：CuGraphBaseConv

一个加速的GraphSAGE层，来自大型图上的归纳表示学习，它利用了cugraph-ops中高度优化的聚合原语：

\[ \begin{align}\begin{aligned}h_{\mathcal{N}(i)}^{(l+1)} &= \mathrm{aggregate} \left(\{h_{j}^{l}, \forall j \in \mathcal{N}(i) \}\right)\\h_{i}^{(l+1)} &= W \cdot \mathrm{concat} (h_{i}^{l}, h_{\mathcal{N}(i)}^{(l+1)})\end{aligned}\end{align} \]

此模块依赖于 pylibcugraphops 包，可以通过 conda install -c nvidia pylibcugraphops=23.04 安装。 pylibcugraphops 23.04 需要 python 3.8.x 或 3.10.x。

注意

这是一个实验性功能。

Parameters:

in_feats (int) – 输入特征大小。
out_feats (int) – 输出特征大小。
aggregator_type (str) – 使用的聚合器类型 (mean, sum, min, max).
feat_drop (float) – 特征上的丢弃率，默认值为：0。
bias (bool) – If True, adds a learnable bias to the output. Default: True.

示例

>>> import dgl
>>> import torch
>>> from dgl.nn import CuGraphSAGEConv
>>> 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 = CuGraphSAGEConv(10, 2, 'mean').to(device)
>>> res = conv(g, feat)
>>> res
tensor([[-1.1690,  0.1952],
        [-1.1690,  0.1952],
        [-1.1690,  0.1952],
        [-1.1690,  0.1952],
        [-1.1690,  0.1952],
        [-1.1690,  0.1952]], device='cuda:0', grad_fn=<AddmmBackward0>)

forward(g, feat, max_in_degree=None)[source]

前向计算。

Parameters:

g (DGLGraph) – The graph.
特征 (torch.Tensor) – 节点特征。形状：\((N, D_{in})\).
max_in_degree (int) – 目标节点的最大入度。只有当 g 是 DGLBlock 时，即二分图时，才有效。当 g 由邻居采样器生成时，该值应设置为相应的 fanout。如果未给出， max_in_degree 将动态计算。

Returns:

输出节点特征。形状：\((N, D_{out})\)。

Return type:

torch.Tensor

reset_parameters()[source]: 重新初始化可学习的参数。